Point seg tree (#51)

* finished refactor

* still has bugs

* saving progress

* nits

* now AC's!!!!!!!!!!!

* now ACs on yosupo

* fix cargo test

* finish docs

* now works with capturing vars

* golf

* saving progress

* fmt

* nit

* fix test

---------

Co-authored-by: Luke Videckis <lukevideckis@gmail.com>

Author: lrvideckis

Cargo.toml

@@ -39,8 +39,8 @@ 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"
+name = "lazy_seg_tree_build_on_array"
+path = "examples/data_structures/lazy_seg_tree_build_on_array.rs"
 
 [[example]]
 name = "trie"
@@ -89,3 +89,11 @@ path = "examples/numbers/primes_is_prime.rs"
 [[example]]
 name = "manacher"
 path = "examples/strings/manacher.rs"
+
+[[example]]
+name = "seg_tree_yosupo"
+path = "examples/data_structures/seg_tree_yosupo.rs"
+
+[[example]]
+name = "seg_tree_aizu"
+path = "examples/data_structures/seg_tree_aizu.rs"

examples/data_structures/lazy_seg_tree.rs

@@ -1,15 +1,15 @@
 // 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;
+use programming_team_code_rust::data_structures::lazy_seg_tree::LazySegTree;
 
 fn main() {
     input! {
         n: usize,
         q: usize
     }
 
-    let mut seg_tree = SegTree::new(n);
+    let mut seg_tree = LazySegTree::new(n);
     for _ in 0..q {
         input! {
             t: usize

…ta_structures/seg_tree_build_on_array.rs …ructures/lazy_seg_tree_build_on_array.rsexamples/data_structures/seg_tree_build_on_array.rs renamed to examples/data_structures/lazy_seg_tree_build_on_array.rs examples/data_structures/seg_tree_build_on_array.rs renamed to examples/data_structures/lazy_seg_tree_build_on_array.rs

@@ -1,7 +1,7 @@
 // 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;
+use programming_team_code_rust::data_structures::lazy_seg_tree::LazySegTree;
 
 fn main() {
     input! {
@@ -10,7 +10,7 @@ fn main() {
         a: [u64; n]
     }
 
-    let mut seg_tree = SegTree::build_on_array(&a);
+    let mut seg_tree = LazySegTree::build_on_array(&a);
     for _ in 0..q {
         input! {
             t: usize

examples/data_structures/seg_tree_aizu.rs

@@ -0,0 +1,33 @@
+// verification-helper: PROBLEM https://judge.u-aizu.ac.jp/onlinejudge/description.jsp?id=DSL_2_B
+
+use proconio::input;
+use programming_team_code_rust::data_structures::seg_tree::SegTree;
+
+fn main() {
+    input! {
+        n: usize,
+        q: usize,
+    }
+    let mut st = SegTree::<usize>::new(n, |x, y| x + y, 0);
+    for _ in 0..q {
+        input! {
+            t: u8,
+        }
+        match t {
+            0 => {
+                input! {
+                    i: usize,
+                    x: usize,
+                }
+                st.set(i - 1, st.query(i - 1..i) + x);
+            }
+            _ => {
+                input! {
+                    le: usize,
+                    ri: usize,
+                }
+                println!("{}", st.query(le - 1..ri));
+            }
+        }
+    }
+}

examples/data_structures/seg_tree_yosupo.rs

@@ -0,0 +1,53 @@
+// verification-helper: PROBLEM https://judge.yosupo.jp/problem/point_set_range_composite
+
+use proconio::input;
+use programming_team_code_rust::data_structures::seg_tree::SegTree;
+
+fn main() {
+    input! {
+        n: usize,
+        q: usize
+    }
+
+    let a = (0..n)
+        .map(|_| {
+            input! {
+                c: usize,
+                d: usize
+            }
+            (c, d)
+        })
+        .collect::<Vec<(usize, usize)>>();
+
+    let mut seg_tree = SegTree::<(usize, usize)>::build_on_array(
+        &a,
+        move |x, y| (x.0 * y.0 % 998244353, (y.0 * x.1 + y.1) % 998244353),
+        (1, 0),
+    );
+
+    for _ in 0..q {
+        input! {
+            t: usize
+        }
+
+        match t {
+            0 => {
+                input! {
+                    idx: usize,
+                    c: usize,
+                    d: usize
+                }
+                seg_tree.set(idx, (c, d));
+            }
+            _ => {
+                input! {
+                    le: usize,
+                    ri: usize,
+                    x: usize
+                }
+                let (c, d) = seg_tree.query(le..ri);
+                println!("{}", (c * x + d) % 998244353);
+            }
+        }
+    }
+}

src/data_structures/lazy_seg_tree.rs

@@ -0,0 +1,135 @@
+//! # Lazy Segment Tree
+
+use std::ops::Range;
+
+/// 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(tr: &Range<usize>) -> usize {
+    let pw2 = 1 << tr.len().ilog2();
+    (tr.start + pw2).min(tr.end - pw2 / 2)
+}
+
+fn op(vl: u64, vr: u64) -> u64 {
+    vl + vr
+}
+
+/// # Example
+/// ```
+/// use programming_team_code_rust::data_structures::lazy_seg_tree::LazySegTree;
+///
+/// let mut seg_tree = LazySegTree::new(10);
+/// seg_tree.update(0..5, 1);
+/// seg_tree.update(5..10, 2);
+/// assert_eq!(seg_tree.query(0..10), 15);
+/// ```
+pub struct LazySegTree {
+    n: usize,
+    tree: Vec<u64>,
+    lazy: Vec<u64>,
+}
+
+impl LazySegTree {
+    /// Creates a new segment tree with n elements
+    ///
+    /// # Complexity
+    /// - Time: O(n)
+    /// - Space: O(n)
+    pub fn new(n: usize) -> Self {
+        LazySegTree {
+            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]);
+        }
+        LazySegTree {
+            n,
+            tree,
+            lazy: vec![0; n],
+        }
+    }
+
+    fn apply(&mut self, delta: u64, tr: &Range<usize>, v: usize) {
+        self.tree[v] += delta * tr.len() as u64;
+        if v < self.n {
+            self.lazy[v] += delta;
+        }
+    }
+
+    fn push(&mut self, tr: &Range<usize>, v: usize) {
+        if self.lazy[v] > 0 {
+            let tm = split(tr);
+            self.apply(self.lazy[v], &(tr.start..tm), 2 * v);
+            self.apply(self.lazy[v], &(tm..tr.end), 2 * v + 1);
+            self.lazy[v] = 0;
+        }
+    }
+
+    /// Updates the range [qr.start, qr.end) 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, qr: Range<usize>, delta: u64) {
+        self.update_impl(&qr, delta, &(0..self.n), 1);
+    }
+
+    fn update_impl(&mut self, qr: &Range<usize>, delta: u64, tr: &Range<usize>, v: usize) {
+        if qr.end <= tr.start || tr.end <= qr.start {
+            return;
+        }
+        if qr.start <= tr.start && tr.end <= qr.end {
+            return self.apply(delta, tr, v);
+        }
+        self.push(tr, v);
+        let tm = split(tr);
+        self.update_impl(qr, delta, &(tr.start..tm), 2 * v);
+        self.update_impl(qr, delta, &(tm..tr.end), 2 * v + 1);
+        self.tree[v] = op(self.tree[2 * v], self.tree[2 * v + 1]);
+    }
+
+    /// Queries the range [qr.start, qr.end)
+    ///
+    /// # Complexity
+    /// - Time: O(log(n))
+    /// - Space: O(log(n)) due to the program stack
+    pub fn query(&mut self, qr: Range<usize>) -> u64 {
+        self.query_impl(&qr, &(0..self.n), 1)
+    }
+
+    fn query_impl(&mut self, qr: &Range<usize>, tr: &Range<usize>, v: usize) -> u64 {
+        if qr.end <= tr.start || tr.end <= qr.start {
+            return 0;
+        }
+        if qr.start <= tr.start && tr.end <= qr.end {
+            return self.tree[v];
+        }
+        self.push(tr, v);
+        let tm = split(tr);
+        op(
+            self.query_impl(qr, &(tr.start..tm), 2 * v),
+            self.query_impl(qr, &(tm..tr.end), 2 * v + 1),
+        )
+    }
+}

src/data_structures/mod.rs

@@ -2,6 +2,7 @@
 pub mod binary_trie;
 pub mod dsu;
 pub mod fenwick;
+pub mod lazy_seg_tree;
 pub mod rmq;
 pub mod seg_tree;
 pub mod trie;