competitive_programming_rs/data_structure/
sparse_table.rs

1pub mod sparse_table {
2    pub struct SparseTable<T, F> {
3        data: Vec<Vec<T>>,
4        op: F,
5    }
6
7    impl<T, F> SparseTable<T, F>
8    where
9        T: Copy,
10        F: Fn(T, T) -> T,
11    {
12        pub fn from(v: &[T], init: T, op: F) -> Self {
13            let size = v.len().next_power_of_two();
14            let count = size.trailing_zeros() as usize + 1;
15            let mut data = vec![vec![init; size]; count];
16            for (i, v) in v.iter().cloned().enumerate() {
17                data[0][i] = v;
18            }
19            for c in 1..count {
20                for i in 0..size {
21                    let next = std::cmp::min(size - 1, i + (1 << (c - 1)));
22                    data[c][i] = op(data[c - 1][i], data[c - 1][next]);
23                }
24            }
25
26            Self { data, op }
27        }
28
29        /// get the result for [l, r)
30        pub fn get(&self, l: usize, r: usize) -> T {
31            assert!(l < r);
32            let length = r - l;
33            if length == 1 {
34                return self.data[0][l];
35            }
36            let block_size = length.next_power_of_two() >> 1;
37            let c = block_size.trailing_zeros() as usize;
38            let left = self.data[c][l];
39            let right = self.data[c][r - block_size];
40            (self.op)(left, right)
41        }
42    }
43}
44
45#[cfg(test)]
46mod test {
47    use super::*;
48    use rand::{thread_rng, Rng};
49    use std::cmp;
50
51    #[test]
52    fn test_small() {
53        let v = vec![1, 9, 9, 1, 0, 7, 1, 7];
54        let sparse_table = sparse_table::SparseTable::from(&v, 0, |a, b| cmp::min(a, b));
55        for l in 0..8 {
56            for r in (l + 1)..9 {
57                let &min = v[l..r].iter().min().unwrap();
58                assert_eq!(sparse_table.get(l, r), min);
59            }
60        }
61    }
62
63    #[test]
64    fn test_sparse_table_with_random_array() {
65        let n = 1000;
66        let v: Vec<u64> = (0..n)
67            .map(|_| thread_rng().gen_range(0, 1000000000))
68            .collect();
69        let sparse_table = sparse_table::SparseTable::from(&v, 0, cmp::min);
70        for l in 0..n {
71            for r in (l + 1)..(n + 1) {
72                let &min = v[l..r].iter().min().unwrap();
73                assert_eq!(sparse_table.get(l, r), min);
74            }
75        }
76    }
77}