haar_lib/ds/
sparse_table_2d.rs1use crate::algebra::traits::*;
3use std::{
4 cmp::{max, min},
5 ops::Range,
6};
7
8pub 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 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 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}