competitive/data_structure/
range_fold_with_upper_bound.rs

1use super::{AbelianGroup, BinaryIndexedTree, RadixSortKey, SliceSortExt};
2use std::ops::Range;
3
4/// Offline range folds over entries whose keys are at most a query bound.
5pub struct RangeFoldWithUpperBound<K, M>
6where
7    M: AbelianGroup,
8{
9    keys: Vec<K>,
10    weights: Vec<M::T>,
11    queries: Vec<(Range<usize>, K)>,
12}
13
14impl<K, M> RangeFoldWithUpperBound<K, M>
15where
16    K: RadixSortKey + Ord,
17    M: AbelianGroup,
18{
19    pub fn new(values: impl IntoIterator<Item = (K, M::T)>) -> Self {
20        let (keys, weights) = values.into_iter().unzip();
21        Self {
22            keys,
23            weights,
24            queries: Vec::new(),
25        }
26    }
27
28    pub fn query(&mut self, range: Range<usize>, upper_bound: K) -> usize {
29        let index = self.queries.len();
30        self.queries.push((range, upper_bound));
31        index
32    }
33
34    pub fn execute(self) -> Vec<M::T> {
35        let mut values: Vec<_> = self
36            .keys
37            .iter()
38            .copied()
39            .enumerate()
40            .map(|(index, key)| (key, index))
41            .collect();
42        values.radix_sort_by_key(|&(key, _)| key);
43        let mut order: Vec<_> = self
44            .queries
45            .iter()
46            .enumerate()
47            .map(|(index, (_, upper_bound))| (*upper_bound, index))
48            .collect();
49        order.radix_sort_by_key(|&(upper_bound, _)| upper_bound);
50
51        let mut bit: BinaryIndexedTree<M> = BinaryIndexedTree::new(values.len());
52        let mut answers = vec![M::unit(); self.queries.len()];
53        let mut inserted = 0;
54        for (upper_bound, index) in order {
55            while inserted < values.len() && values[inserted].0 <= upper_bound {
56                let position = values[inserted].1;
57                bit.update(position, self.weights[position].clone());
58                inserted += 1;
59            }
60            let range = &self.queries[index].0;
61            answers[index] = bit.fold(range.start, range.end);
62        }
63        answers
64    }
65}
66
67#[cfg(test)]
68mod tests {
69    use super::*;
70    use crate::{algebra::AdditiveOperation, rand, tools::Xorshift};
71
72    #[test]
73    fn test_range_fold_with_upper_bound() {
74        let mut rng = Xorshift::default();
75        for _ in 0..100 {
76            rand!(rng, n: 1..200, keys: [-10..10; n], weights: [-100..100; n]);
77            let mut fold: RangeFoldWithUpperBound<_, AdditiveOperation<i64>> =
78                RangeFoldWithUpperBound::new(keys.iter().copied().zip(weights.iter().copied()));
79            let mut expected = Vec::new();
80            for _ in 0..100 {
81                rand!(rng, l: 0..=n, r: l..=n, upper_bound: -10..=10);
82                expected.push(
83                    keys[l..r]
84                        .iter()
85                        .zip(&weights[l..r])
86                        .filter(|&(&key, _)| key <= upper_bound)
87                        .map(|(_, &weight)| weight)
88                        .sum(),
89                );
90                fold.query(l..r, upper_bound);
91            }
92            assert_eq!(fold.execute(), expected);
93        }
94    }
95}
competitive/data_structure/range_fold_with_upper_bound.rs

competitive/data_structure/
range_fold_with_upper_bound.rs
