haar_lib/ds/
segtree_linear_add_range_sum.rs

1//! 区間一次関数加算区間総和セグメントツリー
2
3use crate::math::linear::*;
4use crate::misc::range::range_bounds_to_range;
5use crate::num::one_zero::Zero;
6use crate::trait_alias;
7use std::ops::{Add, AddAssign, Mul, RangeBounds};
8
9trait_alias!(
10    /// [`SegtreeLinearAddRangeSum<T>`]が扱える型
11    Elem: Copy + Zero + Add<Output = Self> + Mul<Output = Self> + AddAssign + PartialEq + From<u32>
12);
13
14/// 区間一次関数加算区間総和セグメントツリー
15pub struct SegtreeLinearAddRangeSum<T> {
16    data: Vec<T>,
17    lazy: Vec<(T, T)>,
18    hsize: usize,
19    original_size: usize,
20}
21
22impl<T: Elem> SegtreeLinearAddRangeSum<T> {
23    /// **Time complexity** $O(n)$
24    ///
25    /// **Space complexity** $O(n)$
26    pub fn new(n: usize) -> Self {
27        let hsize = n.next_power_of_two();
28
29        Self {
30            data: vec![T::zero(); hsize * 2],
31            lazy: vec![(T::zero(), T::zero()); hsize * 2],
32            hsize,
33            original_size: n,
34        }
35    }
36
37    fn _add(a: (T, T), b: (T, T)) -> (T, T) {
38        (a.0 + b.0, a.1 + b.1)
39    }
40
41    fn _propagate(&mut self, i: usize, l: usize, r: usize) {
42        if self.lazy[i] == (T::zero(), T::zero()) {
43            return;
44        }
45        if i < self.hsize {
46            let mut t = self.lazy[i];
47            self.lazy[i << 1] = Self::_add(t, self.lazy[i << 1]);
48            t.0 += t.1 * T::from(((r - l) / 2) as u32);
49            self.lazy[(i << 1) | 1] = Self::_add(t, self.lazy[(i << 1) | 1]);
50        }
51        let len = r - l;
52        let (s, d) = self.lazy[i];
53
54        self.data[i] += s * T::from(len as u32) + d * T::from(((len - 1) * len / 2) as u32);
55        self.lazy[i] = (T::zero(), T::zero());
56    }
57
58    fn _update(&mut self, i: usize, l: usize, r: usize, s: usize, t: usize, a: T, b: T) -> T {
59        self._propagate(i, l, r);
60        if r <= s || t <= l {
61            self.data[i]
62        } else if s <= l && r <= t {
63            self.lazy[i] = Self::_add(self.lazy[i], (a * T::from(l as u32) + b, a));
64            self._propagate(i, l, r);
65            self.data[i]
66        } else {
67            let mid = (l + r) / 2;
68            self.data[i] = self._update(i << 1, l, mid, s, t, a, b)
69                + self._update((i << 1) | 1, mid, r, s, t, a, b);
70            self.data[i]
71        }
72    }
73
74    /// **Time complexity** $O(\log n)$
75    pub fn update(&mut self, range: impl RangeBounds<usize>, linear: Linear<T>) {
76        let (start, end) = range_bounds_to_range(range, 0, self.original_size);
77        self._update(1, 0, self.hsize, start, end, linear.a, linear.b);
78    }
79
80    fn _fold(&mut self, i: usize, l: usize, r: usize, x: usize, y: usize) -> T {
81        self._propagate(i, l, r);
82        if r <= x || y <= l {
83            T::zero()
84        } else if x <= l && r <= y {
85            self.data[i]
86        } else {
87            let mid = (l + r) / 2;
88            self._fold(i << 1, l, mid, x, y) + self._fold((i << 1) | 1, mid, r, x, y)
89        }
90    }
91
92    /// **Time complexity** $O(\log n)$
93    pub fn fold(&mut self, range: impl RangeBounds<usize>) -> T {
94        let (start, end) = range_bounds_to_range(range, 0, self.original_size);
95        self._fold(1, 0, self.hsize, start, end)
96    }
97}
98
99#[cfg(test)]
100mod tests {
101    use super::*;
102    use my_testtools::*;
103    use rand::Rng;
104    use std::ops::Range;
105
106    #[test]
107    fn test() {
108        #![allow(clippy::needless_range_loop)]
109        let mut rng = rand::thread_rng();
110        let n = 100;
111
112        let mut seg = SegtreeLinearAddRangeSum::<i64>::new(n);
113        let mut vec = vec![0; n];
114
115        for _ in 0..300 {
116            let Range { start: l, end: r } = rand_range(&mut rng, 0..n);
117
118            let a = rng.gen_range(0..100);
119            let b = rng.gen_range(0..100);
120
121            seg.update(l..r, Linear { a, b });
122
123            for i in l..r {
124                vec[i] += a * i as i64 + b;
125            }
126
127            let Range { start: l, end: r } = rand_range(&mut rng, 0..n);
128
129            let res = seg.fold(l..r);
130            let ans = vec[l..r].iter().sum::<i64>();
131
132            assert_eq!(res, ans);
133        }
134    }
135}
haar_lib/ds/segtree_linear_add_range_sum.rs

haar_lib/ds/
segtree_linear_add_range_sum.rs
