seg tree (#37)

Author: cameroncuster

Cargo.toml

@@ -30,10 +30,18 @@ path = "examples/data_structures/fenwick_aizu.rs"
 name = "fenwick_yosupo"
 path = "examples/data_structures/fenwick_yosupo.rs"
 
+[[example]]
+name = "lazy_seg_tree"
+path = "examples/data_structures/lazy_seg_tree.rs"
+
 [[example]]
 name = "rmq"
 path = "examples/data_structures/rmq.rs"
 
+[[example]]
+name = "seg_tree_build_on_array"
+path = "examples/data_structures/seg_tree_build_on_array.rs"
+
 [[example]]
 name = "trie"
 path = "examples/data_structures/trie.rs"

examples/data_structures/lazy_seg_tree.rs

@@ -0,0 +1,37 @@
+// verification-helper: PROBLEM https://onlinejudge.u-aizu.ac.jp/courses/library/3/DSL/all/DSL_2_G
+
+use proconio::input;
+use programming_team_code_rust::data_structures::seg_tree::SegTree;
+
+fn main() {
+    input! {
+        n: usize,
+        q: usize
+    }
+
+    let mut seg_tree = SegTree::new(n);
+    for _ in 0..q {
+        input! {
+            t: usize
+        }
+
+        match t {
+            0 => {
+                input! {
+                    le: usize,
+                    ri: usize,
+                    delta: u64
+                }
+                seg_tree.update(le - 1, ri, delta);
+            }
+            1 => {
+                input! {
+                    le: usize,
+                    ri: usize
+                }
+                println!("{}", seg_tree.query(le - 1, ri));
+            }
+            _ => unreachable!(),
+        }
+    }
+}

examples/data_structures/seg_tree_build_on_array.rs

@@ -0,0 +1,37 @@
+// verification-helper: PROBLEM https://judge.yosupo.jp/problem/point_add_range_sum
+
+use proconio::input;
+use programming_team_code_rust::data_structures::seg_tree::SegTree;
+
+fn main() {
+    input! {
+        n: usize,
+        q: usize,
+        a: [u64; n]
+    }
+
+    let mut seg_tree = SegTree::build_on_array(&a);
+    for _ in 0..q {
+        input! {
+            t: usize
+        }
+
+        match t {
+            0 => {
+                input! {
+                    p: usize,
+                    x: u64
+                }
+                seg_tree.update(p, p + 1, x);
+            }
+            1 => {
+                input! {
+                    le: usize,
+                    ri: usize
+                }
+                println!("{}", seg_tree.query(le, ri));
+            }
+            _ => unreachable!(),
+        }
+    }
+}

src/data_structures/mod.rs

@@ -3,4 +3,5 @@ pub mod binary_trie;
 pub mod dsu;
 pub mod fenwick;
 pub mod rmq;
+pub mod seg_tree;
 pub mod trie;

src/data_structures/seg_tree.rs

@@ -0,0 +1,132 @@
+//! # Segment Tree
+
+/// see https://codeforces.com/blog/entry/112755
+///
+/// returns the split point of the range which
+/// makes the segment tree a complete binary tree
+fn split(tl: usize, tr: usize) -> usize {
+    let pw2 = 1 << (tr - tl).ilog2();
+    (tl + pw2).min(tr - pw2 / 2)
+}
+
+fn op(vl: u64, vr: u64) -> u64 {
+    vl + vr
+}
+
+/// # Example
+/// ```
+/// use programming_team_code_rust::data_structures::seg_tree::SegTree;
+///
+/// let mut seg_tree = SegTree::new(10);
+/// seg_tree.update(0, 5, 1);
+/// seg_tree.update(5, 10, 2);
+/// assert_eq!(seg_tree.query(0, 10), 15);
+/// ```
+pub struct SegTree {
+    n: usize,
+    tree: Vec<u64>,
+    lazy: Vec<u64>,
+}
+
+impl SegTree {
+    /// Creates a new segment tree with n elements
+    ///
+    /// # Complexity
+    /// - Time: O(n)
+    /// - Space: O(n)
+    pub fn new(n: usize) -> Self {
+        SegTree {
+            n,
+            tree: vec![0; 2 * n],
+            lazy: vec![0; n],
+        }
+    }
+
+    /// Constructs a segment tree on a given array
+    ///
+    /// # Complexity
+    /// - Time: O(n)
+    /// - Space: O(n)
+    pub fn build_on_array(a: &[u64]) -> Self {
+        let n = a.len();
+        let mut pw2 = 1;
+        while pw2 < n {
+            pw2 *= 2;
+        }
+        let mut tree = vec![0; 2 * n];
+        for i in 0..n {
+            tree[(i + pw2) % n + n] = a[i];
+        }
+        for i in (1..n).rev() {
+            tree[i] = op(tree[2 * i], tree[2 * i + 1]);
+        }
+        SegTree {
+            n,
+            tree,
+            lazy: vec![0; n],
+        }
+    }
+
+    fn apply(&mut self, delta: u64, tl: usize, tr: usize, v: usize) {
+        self.tree[v] += delta * (tr - tl) as u64;
+        if v < self.n {
+            self.lazy[v] += delta;
+        }
+    }
+
+    fn push(&mut self, tl: usize, tm: usize, tr: usize, v: usize) {
+        if self.lazy[v] > 0 {
+            self.apply(self.lazy[v], tl, tm, 2 * v);
+            self.apply(self.lazy[v], tm, tr, 2 * v + 1);
+            self.lazy[v] = 0;
+        }
+    }
+
+    /// Updates the range [le, ri) by adding delta to each element
+    /// in the range
+    ///
+    /// # Complexity
+    /// - Time: O(log(n))
+    /// - Space: O(log(n)) due to the program stack
+    pub fn update(&mut self, le: usize, ri: usize, delta: u64) {
+        self.update_impl(le, ri, delta, 0, self.n, 1);
+    }
+
+    fn update_impl(&mut self, le: usize, ri: usize, delta: u64, tl: usize, tr: usize, v: usize) {
+        if ri <= tl || tr <= le {
+            return;
+        }
+        if le <= tl && tr <= ri {
+            return self.apply(delta, tl, tr, v);
+        }
+        let tm = split(tl, tr);
+        self.push(tl, tm, tr, v);
+        self.update_impl(le, ri, delta, tl, tm, 2 * v);
+        self.update_impl(le, ri, delta, tm, tr, 2 * v + 1);
+        self.tree[v] = op(self.tree[2 * v], self.tree[2 * v + 1]);
+    }
+
+    /// Queries the range [le, ri)
+    ///
+    /// # Complexity
+    /// - Time: O(log(n))
+    /// - Space: O(log(n)) due to the program stack
+    pub fn query(&mut self, le: usize, ri: usize) -> u64 {
+        self.query_impl(le, ri, 0, self.n, 1)
+    }
+
+    fn query_impl(&mut self, le: usize, ri: usize, tl: usize, tr: usize, v: usize) -> u64 {
+        if ri <= tl || tr <= le {
+            return 0;
+        }
+        if le <= tl && tr <= ri {
+            return self.tree[v];
+        }
+        let tm = split(tl, tr);
+        self.push(tl, tm, tr, v);
+        op(
+            self.query_impl(le, ri, tl, tm, 2 * v),
+            self.query_impl(le, ri, tm, tr, 2 * v + 1),
+        )
+    }
+}