haar_lib/ds/
disjoint_sparse_table.rs

1//! 半群の列の区間取得($O(1)$)ができる。
2
3pub use crate::algebra::traits::Semigroup;
4use crate::misc::range::range_bounds_to_range;
5use std::{iter::repeat_n, ops::RangeBounds};
6
7/// 半群の列の区間取得($O(1)$)ができる。
8pub struct DisjointSparseTable<S: Semigroup> {
9    data: Vec<Vec<Option<S>>>,
10    seq: Vec<Option<S>>,
11    size: usize,
12}
13
14impl<S: Semigroup + Clone> DisjointSparseTable<S> {
15    /// 列`seq`から`DisjointSparseTable<S>`を構築する。
16    pub fn new(seq: Vec<S>) -> Self {
17        assert!(!seq.is_empty());
18
19        let size = seq.len();
20        let log_size = usize::BITS as usize - (size - 1).leading_zeros() as usize;
21        let mut data = vec![vec![None; 1 << log_size]; log_size];
22
23        let seq = seq
24            .into_iter()
25            .map(Some)
26            .chain(repeat_n(None, (1 << log_size) - size))
27            .collect::<Vec<_>>();
28
29        for (i, x) in seq.iter().enumerate() {
30            data[0][i] = x.clone();
31        }
32
33        let mut this = Self { data, seq, size };
34        this.build(0, 1 << log_size, log_size - 1);
35
36        this
37    }
38
39    fn build(&mut self, l: usize, r: usize, d: usize) {
40        let m = (l + r) / 2;
41
42        self.data[d][m] = self.seq[m].clone();
43        for i in m + 1..r {
44            self.data[d][i] = match (self.data[d][i - 1].clone(), self.seq[i].clone()) {
45                (Some(x), Some(y)) => Some(x.op(y)),
46                (a, None) => a,
47                (None, a) => a,
48            }
49        }
50
51        self.data[d][m - 1] = self.seq[m - 1].clone();
52        for i in (l..m - 1).rev() {
53            self.data[d][i] = match (self.seq[i].clone(), self.data[d][i + 1].clone()) {
54                (Some(x), Some(y)) => Some(x.op(y)),
55                (a, None) => a,
56                (None, a) => a,
57            }
58        }
59
60        if d > 0 {
61            self.build(l, m, d - 1);
62            self.build(m, r, d - 1);
63        }
64    }
65
66    /// **Time complexity** $O(1)$
67    pub fn fold(&self, range: impl RangeBounds<usize>) -> Option<S> {
68        let (l, r) = range_bounds_to_range(range, 0, self.size);
69
70        if l == r {
71            None
72        } else {
73            let r = r - 1;
74
75            if l == r {
76                self.seq[l].clone()
77            } else {
78                let k = usize::BITS as usize - 1 - (l ^ r).leading_zeros() as usize;
79                match (self.data[k][l].clone(), self.data[k][r].clone()) {
80                    (Some(x), Some(y)) => Some(x.op(y)),
81                    (a, None) => a,
82                    (None, a) => a,
83                }
84            }
85        }
86    }
87}
88
89#[cfg(test)]
90mod tests {
91    use super::*;
92    use crate::algebra::sum::*;
93    use my_testtools::*;
94    use rand::Rng;
95
96    #[test]
97    fn test() {
98        let mut rng = rand::thread_rng();
99
100        let n = 100;
101        let a = (0..n)
102            .map(|_| Sum(rng.gen::<u32>() % 10000))
103            .collect::<Vec<_>>();
104        let s = DisjointSparseTable::<Sum<u32>>::new(a.clone());
105
106        for _ in 0..100 {
107            let lr = rand_range(&mut rng, 0..n);
108
109            assert_eq!(
110                s.fold(lr.clone()),
111                if lr.is_empty() {
112                    None
113                } else {
114                    Some(a[lr].iter().cloned().fold_m())
115                }
116            );
117        }
118    }
119}