ranges (#44)

* ranges

* lazy ranges

Author: cameroncuster

examples/data_structures/lazy_seg_tree.rs

@@ -22,14 +22,14 @@ fn main() {
                     ri: usize,
                     delta: u64
                 }
-                seg_tree.update(le - 1, ri, delta);
+                seg_tree.update(le - 1..ri, delta);
             }
             _ => {
                 input! {
                     le: usize,
                     ri: usize
                 }
-                println!("{}", seg_tree.query(le - 1, ri));
+                println!("{}", seg_tree.query(le - 1..ri));
             }
         }
     }

examples/data_structures/seg_tree_build_on_array.rs

@@ -22,14 +22,14 @@ fn main() {
                     p: usize,
                     x: u64
                 }
-                seg_tree.update(p, p + 1, x);
+                seg_tree.update(p..p + 1, x);
             }
             _ => {
                 input! {
                     le: usize,
                     ri: usize
                 }
-                println!("{}", seg_tree.query(le, ri));
+                println!("{}", seg_tree.query(le..ri));
             }
         }
     }

src/data_structures/seg_tree.rs

@@ -1,12 +1,14 @@
 //! # 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(tl: usize, tr: usize) -> usize {
-    let pw2 = 1 << (tr - tl).ilog2();
-    (tl + pw2).min(tr - pw2 / 2)
+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 {
@@ -18,9 +20,9 @@ fn op(vl: u64, vr: u64) -> u64 {
 /// 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);
+/// 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,
@@ -67,17 +69,18 @@ impl SegTree {
         }
     }
 
-    fn apply(&mut self, delta: u64, tl: usize, tr: usize, v: usize) {
-        self.tree[v] += delta * (tr - tl) as u64;
+    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, tl: usize, tm: usize, tr: usize, v: usize) {
+    fn push(&mut self, tr: &Range<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);
+            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;
         }
     }
@@ -88,21 +91,21 @@ impl SegTree {
     /// # 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);
+    pub fn update(&mut self, qr: Range<usize>, delta: u64) {
+        self.update_impl(&qr, 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 {
+    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 le <= tl && tr <= ri {
-            return self.apply(delta, tl, tr, v);
+        if qr.start <= tr.start && tr.end <= qr.end {
+            return self.apply(delta, 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);
+        let tm = split(tr);
+        self.push(tr, v);
+        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]);
     }
 
@@ -111,22 +114,22 @@ impl SegTree {
     /// # 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)
+    pub fn query(&mut self, qr: Range<usize>) -> u64 {
+        self.query_impl(&qr, &(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 {
+    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 le <= tl && tr <= ri {
+        if qr.start <= tr.start && tr.end <= qr.end {
             return self.tree[v];
         }
-        let tm = split(tl, tr);
-        self.push(tl, tm, tr, v);
+        let tm = split(tr);
+        self.push(tr, v);
         op(
-            self.query_impl(le, ri, tl, tm, 2 * v),
-            self.query_impl(le, ri, tm, tr, 2 * v + 1),
+            self.query_impl(qr, &(tr.start..tm), 2 * v),
+            self.query_impl(qr, &(tm..tr.end), 2 * v + 1),
         )
     }
 }