contest_algorithms/range_query/
sqrt_decomp.rs

1/// A generic implementation of Mo's algorithm, aka Query Sqrt Decomposition.
2/// It answers q offline queries over intervals in 0..n by shifting the query
3/// interval's endpoints by one position at a time.
4/// Each endpoint incurs a total cost of at most n * sqrt(q * L_OP * R_OP).
5pub trait MoState {
6    type Q;
7    type A;
8
9    /// cost ratio L_OP / R_OP between a left endpoint and a right endpoint move
10    const L_R_RATIO: f64 = 1.0;
11
12    fn query(&self, q: &Self::Q) -> Self::A;
13    fn insert_left(&mut self, pos: usize);
14    fn remove_left(&mut self, pos: usize);
15
16    fn insert_right(&mut self, pos: usize) {
17        self.insert_left(pos);
18    }
19    fn remove_right(&mut self, pos: usize) {
20        self.remove_left(pos);
21    }
22    /// After initializing self to a state corresponding to an empty interval,
23    /// call this function to answer all your queries.
24    fn process(&mut self, queries: &[(usize, usize, Self::Q)]) -> Vec<Self::A> {
25        let q = queries.len();
26        let mut q_positions: Vec<usize> = (0..q).collect();
27        if let Some(max_r) = queries.iter().map(|&(_, r, _)| r).max() {
28            let q_adjusted = q as f64 * Self::L_R_RATIO;
29            let bucket_width = 1 + max_r / q_adjusted.sqrt() as usize;
30            q_positions.sort_unstable_by_key(|&i| {
31                let (l, mut r) = (queries[i].0, queries[i].1);
32                let bucket = l / bucket_width;
33                if bucket % 2 != 0 {
34                    r = max_r - r;
35                }
36                (bucket, r)
37            });
38        }
39
40        let (mut cur_l, mut cur_r) = (1, 0);
41        let mut answers = Vec::with_capacity(queries.len());
42        for i in q_positions {
43            let (l, r, ref q) = queries[i];
44            while cur_l > l {
45                cur_l -= 1;
46                self.insert_left(cur_l);
47            }
48            while cur_r < r {
49                cur_r += 1;
50                self.insert_right(cur_r);
51            }
52            while cur_l < l {
53                self.remove_left(cur_l);
54                cur_l += 1;
55            }
56            while cur_r > r {
57                self.remove_right(cur_r);
58                cur_r -= 1;
59            }
60            answers.push((i, self.query(q)));
61        }
62        answers.sort_unstable_by_key(|&(i, _)| i);
63        answers.into_iter().map(|(_, ans)| ans).collect()
64    }
65}
66
67pub struct DistinctVals {
68    vals: Vec<usize>,
69    counts: Vec<usize>,
70    distinct: usize,
71}
72impl DistinctVals {
73    pub fn new(vals: Vec<usize>) -> Self {
74        let &max_val = vals.iter().max().unwrap_or(&0);
75        Self {
76            vals,
77            counts: vec![0; max_val + 1],
78            distinct: 0,
79        }
80    }
81}
82impl MoState for DistinctVals {
83    type Q = ();
84    type A = usize;
85    fn query(&self, _: &Self::Q) -> Self::A {
86        self.distinct
87    }
88    fn insert_left(&mut self, pos: usize) {
89        let v = self.vals[pos];
90        if self.counts[v] == 0 {
91            self.distinct += 1;
92        }
93        self.counts[v] += 1;
94    }
95    fn remove_left(&mut self, pos: usize) {
96        let v = self.vals[pos];
97        self.counts[v] -= 1;
98        if self.counts[v] == 0 {
99            self.distinct -= 1;
100        }
101    }
102}
103
104#[cfg(test)]
105mod test {
106    use super::*;
107
108    #[test]
109    fn test_mos_algorithm() {
110        let queries = vec![(0, 2, ()), (5, 5, ()), (2, 6, ()), (0, 6, ())];
111        let arr = vec![4, 8, 4, 7, 1, 9, 8];
112
113        let answers = DistinctVals::new(arr).process(&queries);
114
115        assert_eq!(answers, vec![2, 1, 5, 5]);
116    }
117}
contest_algorithms/range_query/sqrt_decomp.rs

contest_algorithms/range_query/
sqrt_decomp.rs
