haar_lib/ds/
segtree_2d.rs

1//! 二次元のセグメント木
2use crate::algebra::traits::*;
3use std::ops::Range;
4
5/// 二次元のセグメント木
6pub struct Segtree2D<M: Monoid + Commutative> {
7    monoid: M,
8    data: Vec<Vec<M::Element>>,
9    w: usize,
10    h: usize,
11}
12
13impl<M: Monoid + Commutative> Segtree2D<M>
14where
15    M::Element: Clone,
16{
17    /// **Time complexity** $O(wh)$
18    ///
19    /// **Space complexity** $O(wh)$
20    pub fn new(monoid: M, w: usize, h: usize) -> Self {
21        let w = w.next_power_of_two() * 2;
22        let h = h.next_power_of_two() * 2;
23        let data = vec![vec![monoid.id(); h]; w];
24        Self { monoid, data, w, h }
25    }
26
27    fn __fold(&self, l: usize, r: usize, x: usize) -> M::Element {
28        let mut l = l + self.h / 2;
29        let mut r = r + self.h / 2;
30
31        let mut ret = self.monoid.id();
32        let a = &self.data[x];
33
34        while l < r {
35            if r & 1 == 1 {
36                r -= 1;
37                ret = self.monoid.op(ret, a[r].clone());
38            }
39            if l & 1 == 1 {
40                ret = self.monoid.op(ret, a[l].clone());
41                l += 1;
42            }
43            l >>= 1;
44            r >>= 1;
45        }
46
47        ret
48    }
49
50    /// **Time complexity** $O(\log w \log h)$
51    pub fn fold_2d(
52        &self,
53        Range { start: x1, end: x2 }: Range<usize>,
54        Range { start: y1, end: y2 }: Range<usize>,
55    ) -> M::Element {
56        let mut l = x1 + self.w / 2;
57        let mut r = x2 + self.w / 2;
58
59        let mut ret = self.monoid.id();
60
61        while l < r {
62            if r & 1 == 1 {
63                r -= 1;
64                ret = self.monoid.op(ret, self.__fold(y1, y2, r));
65            }
66            if l & 1 == 1 {
67                ret = self.monoid.op(ret, self.__fold(y1, y2, l));
68                l += 1;
69            }
70            l >>= 1;
71            r >>= 1;
72        }
73
74        ret
75    }
76
77    /// **Time complexity** $O(1)$
78    pub fn get(&self, i: usize, j: usize) -> M::Element {
79        self.data[i + self.w / 2][j + self.h / 2].clone()
80    }
81
82    /// **Time complexity** $O(\log w \log h)$
83    pub fn assign(&mut self, i: usize, j: usize, value: M::Element) {
84        let i = i + self.w / 2;
85        let j = j + self.h / 2;
86
87        self.data[i][j] = value;
88
89        let mut x = i >> 1;
90        while x > 0 {
91            self.data[x][j] = self.monoid.op(
92                self.data[x << 1][j].clone(),
93                self.data[(x << 1) | 1][j].clone(),
94            );
95            x >>= 1;
96        }
97
98        let mut y = j >> 1;
99        while y > 0 {
100            let mut x = i;
101            while x > 0 {
102                self.data[x][y] = self.monoid.op(
103                    self.data[x][y << 1].clone(),
104                    self.data[x][(y << 1) | 1].clone(),
105                );
106                x >>= 1;
107            }
108            y >>= 1;
109        }
110    }
111
112    /// **Time complexity** $O(\log w \log h)$
113    pub fn update(&mut self, i: usize, j: usize, value: M::Element) {
114        let value = self.monoid.op(value, self.get(i, j));
115        self.assign(i, j, value);
116    }
117}
118
119#[cfg(test)]
120mod tests {
121    use super::*;
122    use crate::algebra::sum::*;
123    use my_testtools::*;
124    use rand::Rng;
125
126    #[test]
127    fn test() {
128        #![allow(clippy::needless_range_loop)]
129        let w = 300;
130        let h = 100;
131
132        let mut rng = rand::thread_rng();
133
134        let m = Sum::<u64>::new();
135        let mut seg = Segtree2D::new(m, w, h);
136        let mut a = vec![vec![m.id(); h]; w];
137
138        for i in 0..w {
139            for j in 0..h {
140                let x = rng.gen::<u64>() % 10000;
141
142                a[i][j] = x;
143                seg.assign(i, j, x);
144            }
145        }
146
147        for _ in 0..100 {
148            let i = rng.gen::<usize>() % w;
149            let j = rng.gen::<usize>() % h;
150            let x = rng.gen::<u64>() % 10000;
151
152            seg.assign(i, j, x);
153            a[i][j] = x;
154
155            let wr = rand_range(&mut rng, 0..w);
156            let hr = rand_range(&mut rng, 0..h);
157
158            let res = seg.fold_2d(wr.clone(), hr.clone());
159
160            let ans = a[wr]
161                .iter()
162                .map(|a| a[hr.clone()].iter().cloned().fold_m(&m))
163                .fold_m(&m);
164
165            assert_eq!(res, ans);
166        }
167    }
168}
haar_lib/ds/segtree_2d.rs

haar_lib/ds/
segtree_2d.rs
