competitive/data_structure/
range_frequency.rs

1use super::{AdditiveOperation, BinaryIndexedTree};
2use std::{collections::HashMap, hash::Hash, mem::replace};
3
4#[derive(Debug, Clone)]
5enum RangeFrequencyQuery {
6    Add {
7        index: usize,
8    },
9    Remove {
10        index: usize,
11    },
12    Query {
13        left: usize,
14        right: usize,
15        output_index: usize,
16    },
17}
18
19#[derive(Debug, Clone)]
20pub struct RangeFrequency<T>
21where
22    T: Clone + Eq + Hash,
23{
24    array: Vec<T>,
25    queries: HashMap<T, Vec<RangeFrequencyQuery>>,
26    output_size: usize,
27}
28
29impl<T> RangeFrequency<T>
30where
31    T: Clone + Eq + Hash,
32{
33    pub fn new(array: Vec<T>) -> Self {
34        let mut queries = HashMap::<T, Vec<RangeFrequencyQuery>>::new();
35        for (index, value) in array.iter().cloned().enumerate() {
36            queries
37                .entry(value)
38                .or_default()
39                .push(RangeFrequencyQuery::Add { index });
40        }
41        Self {
42            array,
43            queries,
44            output_size: 0,
45        }
46    }
47
48    pub fn set(&mut self, index: usize, value: T) {
49        let old_value = replace(&mut self.array[index], value);
50        self.queries
51            .entry(old_value)
52            .or_default()
53            .push(RangeFrequencyQuery::Remove { index });
54        self.queries
55            .entry(self.array[index].clone())
56            .or_default()
57            .push(RangeFrequencyQuery::Add { index });
58    }
59
60    pub fn query(&mut self, left: usize, right: usize, value: T) -> usize {
61        let output_index = self.output_size;
62        self.queries
63            .entry(value)
64            .or_default()
65            .push(RangeFrequencyQuery::Query {
66                left,
67                right,
68                output_index,
69            });
70        self.output_size += 1;
71        output_index
72    }
73
74    pub fn execute_with_callback(mut self, mut callback: impl FnMut(usize, usize)) {
75        let mut processor = RangeFrequencyProcessor::new(self.array.len());
76        for (index, value) in self.array.into_iter().enumerate() {
77            self.queries
78                .entry(value)
79                .or_default()
80                .push(RangeFrequencyQuery::Remove { index });
81        }
82        for (_, queries) in self.queries {
83            for query in queries {
84                match query {
85                    RangeFrequencyQuery::Add { index } => {
86                        processor.add(index);
87                    }
88                    RangeFrequencyQuery::Remove { index } => {
89                        processor.remove(index);
90                    }
91                    RangeFrequencyQuery::Query {
92                        left,
93                        right,
94                        output_index,
95                    } => {
96                        callback(output_index, processor.query(left, right));
97                    }
98                }
99            }
100        }
101    }
102
103    pub fn execute(self) -> Vec<usize> {
104        let mut results = vec![0; self.output_size];
105        self.execute_with_callback(|i, v| results[i] = v);
106        results
107    }
108}
109
110#[derive(Debug, Clone)]
111struct RangeFrequencyProcessor {
112    bit: BinaryIndexedTree<AdditiveOperation<i32>>,
113    data: Vec<u64>,
114}
115
116impl RangeFrequencyProcessor {
117    fn new(size: usize) -> Self {
118        Self {
119            bit: BinaryIndexedTree::new(size.div_ceil(64)),
120            data: vec![0; size.div_ceil(64)],
121        }
122    }
123
124    fn add(&mut self, index: usize) {
125        let (block, bit) = (index / 64, index % 64);
126        assert!(self.data[block] & (1 << bit) == 0);
127        self.data[block] |= 1 << bit;
128        self.bit.update(block, 1);
129    }
130
131    fn remove(&mut self, index: usize) {
132        let (i, j) = (index / 64, index % 64);
133        assert!(self.data[i] & (1 << j) != 0);
134        self.data[i] &= !(1 << j);
135        self.bit.update(i, -1);
136    }
137
138    fn query(&self, left: usize, right: usize) -> usize {
139        if left >= right {
140            return 0;
141        }
142        let right = right - 1;
143        let (li, lj) = (left / 64, left % 64);
144        let (ri, rj) = (right / 64, right % 64);
145        let rj_r = 63 - rj;
146        if li == ri {
147            (self.data[li] << rj_r >> (lj + rj_r)).count_ones() as usize
148        } else {
149            let mut ans = self.bit.fold(li + 1, ri) as usize;
150            ans += (self.data[li] >> lj).count_ones() as usize;
151            ans += (self.data[ri] << rj_r).count_ones() as usize;
152            ans
153        }
154    }
155}
156
157#[cfg(test)]
158mod tests {
159    use super::*;
160    use crate::{
161        rand,
162        tools::{NotEmptySegment as Nes, Xorshift},
163    };
164
165    #[test]
166    fn test_range_frequency() {
167        let mut rng = Xorshift::default();
168        for _ in 0..100 {
169            rand!(rng, n: 1..200, mut a: [0..10; n]);
170            let mut rf = RangeFrequency::new(a.clone());
171            for _ in 0..100 {
172                rand!(rng, i: 0..n, v: 0..10);
173                rf.set(i, v);
174                a[i] = v;
175            }
176            let mut expected = vec![];
177            for _ in 0..100 {
178                let (l, r) = rng.random(Nes(n));
179                for v in 0..10 {
180                    expected.push(a[l..r].iter().filter(|&&x| x == v).count());
181                    rf.query(l, r, v);
182                }
183            }
184            let result = rf.execute();
185            assert_eq!(result, expected);
186        }
187    }
188}
competitive/data_structure/range_frequency.rs

competitive/data_structure/
range_frequency.rs
