Merge pull request #152 from ia7ck/segtree-iter

[segtree] 非再帰の実装にする

Author: ia7ck

algo/segment_tree/examples/point_set_range_composite.rs

@@ -16,7 +16,7 @@ fn main() {
         (c * a % mo, (c * b % mo + d) % mo)
     });
     for (i, &(a, b)) in ab.iter().enumerate() {
-        seg.update(i, (a, b));
+        seg.set(i, (a, b));
     }
     for _ in 0..q {
         input! {
@@ -28,7 +28,7 @@ fn main() {
                 c: u64,
                 d: u64,
             }
-            seg.update(p, (c, d));
+            seg.set(p, (c, d));
         } else {
             input! {
                 l: usize,

algo/segment_tree/src/lib.rs

@@ -1,17 +1,19 @@
 use std::fmt;
-use std::ops::{Bound, Range, RangeBounds};
+use std::ops::{Bound, RangeBounds};
 
 /// __注意⚠__ この実装は遅いので time limit の厳しい問題には代わりに ACL のセグメントツリーを使うこと。
 ///
 /// セグメントツリーです。
 #[derive(Clone)]
 pub struct SegmentTree<T, F> {
+    original_n: usize,
     n: usize,
     dat: Vec<T>,
     e: T,
     multiply: F,
 }
 
+// https://hcpc-hokudai.github.io/archive/structure_segtree_001.pdf
 impl<T, F> SegmentTree<T, F>
 where
     T: Clone,
@@ -21,33 +23,43 @@ where
     ///
     /// `multiply` は fold に使う二項演算です。
     pub fn new(n: usize, e: T, multiply: F) -> Self {
+        let original_n = n;
         let n = n.next_power_of_two();
         Self {
+            original_n,
             n,
-            dat: vec![e.clone(); n * 2 - 1],
+            dat: vec![e.clone(); n * 2], // dat[0] is unused
             e,
             multiply,
         }
     }
 
     /// 列の `i` 番目の要素を取得します。
     pub fn get(&self, i: usize) -> &T {
-        &self.dat[i + self.n - 1]
+        assert!(i < self.original_n);
+        &self.dat[i + self.n]
     }
 
     /// 列の `i` 番目の要素を `x` で更新します。
-    pub fn update(&mut self, i: usize, x: T) {
-        let mut k = i + self.n - 1;
-        self.dat[k] = x;
-        while k > 0 {
-            k = (k - 1) / 2;
-            self.dat[k] = (self.multiply)(&self.dat[k * 2 + 1], &self.dat[k * 2 + 2]);
+    pub fn set(&mut self, i: usize, x: T) {
+        self.update(i, |_| x);
+    }
+
+    /// 列の `i` 番目の要素を `f` で更新します。
+    pub fn update<U>(&mut self, i: usize, f: U)
+    where
+        U: FnOnce(&T) -> T,
+    {
+        assert!(i < self.original_n);
+        let mut k = i + self.n;
+        self.dat[k] = f(&self.dat[k]);
+        while k > 1 {
+            k >>= 1;
+            self.dat[k] = (self.multiply)(&self.dat[k << 1], &self.dat[k << 1 | 1]);
         }
     }
 
     /// `range` が `l..r` として、`multiply(l番目の要素, multiply(..., multiply(r-2番目の要素, r-1番目の要素)))` の値を返します。
-    ///
-    /// 実際のアルゴリズムは、結合法則を使って `1 + (2 + (3 + 4))` ではなく `(1 + 2) + (3 + 4)` のように計算しています。
     pub fn fold(&self, range: impl RangeBounds<usize>) -> T {
         let start = match range.start_bound() {
             Bound::Included(&start) => start,
@@ -59,21 +71,31 @@ where
             Bound::Excluded(&end) => end,
             Bound::Unbounded => self.n,
         };
-        assert!(end <= self.n);
-        self._fold(&(start..end), 0, 0..self.n)
+        assert!(start <= end && end <= self.original_n);
+        self._fold(start, end)
     }
-    fn _fold(&self, range: &Range<usize>, i: usize, i_range: Range<usize>) -> T {
-        if range.end <= i_range.start || i_range.end <= range.start {
-            return self.e.clone();
-        }
-        if range.start <= i_range.start && i_range.end <= range.end {
-            return self.dat[i].clone();
+
+    fn _fold(&self, mut l: usize, mut r: usize) -> T {
+        let mut acc_l = self.e.clone();
+        let mut acc_r = self.e.clone();
+        l += self.n;
+        r += self.n;
+        while l < r {
+            if l & 1 == 1 {
+                // 右の子だったらいま足しておかないといけない
+                // 左の子だったら祖先のどれかで足されるのでよい
+                acc_l = (self.multiply)(&acc_l, &self.dat[l]);
+                l += 1;
+            }
+            if r & 1 == 1 {
+                // r が exclusive であることに注意する
+                r -= 1;
+                acc_r = (self.multiply)(&self.dat[r], &acc_r);
+            }
+            l >>= 1;
+            r >>= 1;
         }
-        let m = (i_range.start + i_range.end) / 2;
-        (self.multiply)(
-            &self._fold(range, i * 2 + 1, i_range.start..m),
-            &self._fold(range, i * 2 + 2, m..i_range.end),
-        )
+        (self.multiply)(&acc_l, &acc_r)
     }
 }
 
@@ -82,7 +104,7 @@ where
     T: fmt::Debug,
 {
     fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
-        write!(f, "{:?}", &self.dat[(self.n - 1)..])
+        write!(f, "{:?}", &self.dat[self.n..])
     }
 }
 
@@ -95,7 +117,7 @@ mod tests {
         let s = "abcdefgh";
         let mut seg = SegmentTree::new(s.len(), String::new(), |a, b| format!("{a}{b}"));
         for (i, c) in s.chars().enumerate() {
-            seg.update(i, c.to_string());
+            seg.set(i, c.to_string());
         }
 
         for i in 0..s.len() {
@@ -117,7 +139,7 @@ mod tests {
     fn single_element() {
         let mut seg = SegmentTree::new(1, 0, |a, b| a + b);
         assert_eq!(seg.get(0), &0);
-        seg.update(0, 42);
+        seg.set(0, 42);
         assert_eq!(seg.get(0), &42);
     }
 }