haar_lib/ds/
sparse_table_2d.rs

1//! 冪等性と結合性をもつ2次元列の区間取得($O(1)$)ができる。
2use crate::algebra::traits::*;
3use std::{
4    cmp::{max, min},
5    ops::Range,
6};
7
8/// 冪等性と結合性をもつ2次元列の区間取得($O(1)$)ができる。
9pub struct SparseTable2D<A: Semigroup + Idempotence> {
10    data: Vec<Vec<Vec<Vec<A>>>>,
11    log_table: Vec<usize>,
12}
13
14impl<A: Semigroup + Idempotence + Clone + Default> SparseTable2D<A> {
15    /// **Time complexity** $O(nm \log n \log m)$
16    ///
17    /// **Space complexity** $O(nm \log n \log m)$
18    pub fn new(s: Vec<Vec<A>>) -> Self {
19        let n = s.len();
20        let m = s[0].len();
21        let logn = n.next_power_of_two().trailing_zeros() as usize + 1;
22        let logm = m.next_power_of_two().trailing_zeros() as usize + 1;
23
24        let mut data = vec![vec![vec![vec![A::default(); logm]; m]; logn]; n];
25
26        for i in 0..n {
27            for j in 0..m {
28                data[i][0][j][0] = s[i][j].clone();
29            }
30
31            for y in 1..logm {
32                for j in 0..m {
33                    data[i][0][j][y] = A::op(
34                        data[i][0][j][y - 1].clone(),
35                        data[i][0][min(m - 1, j + (1 << (y - 1)))][y - 1].clone(),
36                    );
37                }
38            }
39        }
40
41        for x in 1..logn {
42            for i in 0..n {
43                for y in 0..logm {
44                    for j in 0..m {
45                        data[i][x][j][y] = A::op(
46                            data[i][x - 1][j][y].clone(),
47                            data[min(n - 1, i + (1 << (x - 1)))][x - 1][j][y].clone(),
48                        );
49                    }
50                }
51            }
52        }
53
54        let mut log_table = vec![0; max(n, m) + 1];
55        for i in 2..=max(n, m) {
56            log_table[i] = log_table[i >> 1] + 1;
57        }
58
59        Self { data, log_table }
60    }
61
62    /// **Time complexity** $O(1)$
63    pub fn fold_2d(
64        &self,
65        Range { start: r1, end: r2 }: Range<usize>,
66        Range { start: c1, end: c2 }: Range<usize>,
67    ) -> Option<A> {
68        if r1 == r2 || c1 == c2 {
69            return None;
70        }
71        let kr = self.log_table[r2 - r1];
72        let kc = self.log_table[c2 - c1];
73
74        let x = A::op(
75            self.data[r1][kr][c1][kc].clone(),
76            self.data[r1][kr][c2 - (1 << kc)][kc].clone(),
77        );
78        let y = A::op(
79            self.data[r2 - (1 << kr)][kr][c1][kc].clone(),
80            self.data[r2 - (1 << kr)][kr][c2 - (1 << kc)][kc].clone(),
81        );
82
83        Some(A::op(x, y))
84    }
85}
86
87#[cfg(test)]
88mod tests {
89    use super::*;
90    use crate::{algebra::min_max::Max, iter::collect::CollectVec};
91    use rand::Rng;
92    use std::fmt::Debug;
93
94    fn test<A>(s: Vec<Vec<A>>)
95    where
96        A: Semigroup + Idempotence + Identity + Copy + Default + PartialEq + Debug,
97    {
98        let st = SparseTable2D::new(s.clone());
99        let n = s.len();
100        let m = s[0].len();
101
102        for x1 in 0..n {
103            for x2 in x1..=n {
104                for y1 in 0..m {
105                    for y2 in y1..=m {
106                        let ans = &s[x1..x2]
107                            .iter()
108                            .map(|v| v[y1..y2].iter().cloned().fold_m())
109                            .fold_m();
110
111                        assert_eq!(*ans, st.fold_2d(x1..x2, y1..y2).unwrap_or(A::id()));
112                    }
113                }
114            }
115        }
116    }
117
118    #[test]
119    fn test_max() {
120        let mut rng = rand::thread_rng();
121        let n = 30;
122        let m = 30;
123        let s = (0..n)
124            .map(|_| (0..m).map(|_| Max(rng.gen::<u64>())).collect_vec())
125            .collect_vec();
126        test(s);
127    }
128}