competitive/data_structure/
range_frequency.rs1use super::{AdditiveOperation, BinaryIndexedTree, FibHashMap};
2use std::{collections::hash_map::Entry, hash::Hash, mem::replace};
3
4#[derive(Debug, Clone, Copy)]
5enum RangeFrequencyQuery {
6 Add {
7 index: u32,
8 },
9 Remove {
10 index: u32,
11 },
12 Query {
13 left: u32,
14 right: u32,
15 output_index: u32,
16 },
17}
18
19#[derive(Debug, Clone)]
20pub struct RangeFrequency<T>
21where
22 T: Clone + Eq + Hash,
23{
24 array: Vec<u32>,
25 values: FibHashMap<T, u32>,
26 events: Vec<(u32, RangeFrequencyQuery)>,
27 queried: Vec<u8>,
28 static_queries: Option<Vec<(u32, u32, u32, u32)>>,
29 zero_queries: Vec<u32>,
30 output_size: usize,
31}
32
33impl<T> RangeFrequency<T>
34where
35 T: Clone + Eq + Hash,
36{
37 pub fn new(array: Vec<T>) -> Self {
38 let mut result = Self {
39 array: Vec::with_capacity(array.len()),
40 values: FibHashMap::with_capacity_and_hasher(array.len(), Default::default()),
41 events: Vec::new(),
42 queried: Vec::new(),
43 static_queries: Some(Vec::new()),
44 zero_queries: Vec::new(),
45 output_size: 0,
46 };
47 for value in array {
48 let value = result.value_id(value);
49 result.array.push(value);
50 }
51 result
52 }
53
54 fn value_id(&mut self, value: T) -> u32 {
55 match self.values.entry(value) {
56 Entry::Occupied(entry) => *entry.get(),
57 Entry::Vacant(entry) => {
58 let id = self.queried.len() as u32;
59 entry.insert(id);
60 self.queried.push(0);
61 id
62 }
63 }
64 }
65
66 pub fn set(&mut self, index: usize, value: T) {
67 if let Some(queries) = self.static_queries.take() {
68 self.events.reserve(self.array.len() + queries.len() + 2);
69 for (index, &value) in self.array.iter().enumerate() {
70 self.events.push((
71 value,
72 RangeFrequencyQuery::Add {
73 index: index as u32,
74 },
75 ));
76 }
77 for (left, right, value, output_index) in queries {
78 self.events.push((
79 value,
80 RangeFrequencyQuery::Query {
81 left,
82 right,
83 output_index,
84 },
85 ));
86 }
87 }
88 let value = self.value_id(value);
89 let old_value = replace(&mut self.array[index], value);
90 self.events.push((
91 old_value,
92 RangeFrequencyQuery::Remove {
93 index: index as u32,
94 },
95 ));
96 self.events.push((
97 value,
98 RangeFrequencyQuery::Add {
99 index: index as u32,
100 },
101 ));
102 }
103
104 pub fn query(&mut self, left: usize, right: usize, value: T) -> usize {
105 let output_index = self.output_size;
106 if let Some(&value) = self.values.get(&value) {
107 self.queried[value as usize] = 1;
108 if let Some(queries) = &mut self.static_queries {
109 queries.push((left as u32, right as u32, value, output_index as u32));
110 } else {
111 self.events.push((
112 value,
113 RangeFrequencyQuery::Query {
114 left: left as u32,
115 right: right as u32,
116 output_index: output_index as u32,
117 },
118 ));
119 }
120 } else {
121 self.zero_queries.push(output_index as u32);
122 }
123 self.output_size += 1;
124 output_index
125 }
126
127 pub fn execute_with_callback(mut self, mut callback: impl FnMut(usize, usize)) {
128 for output_index in self.zero_queries {
129 callback(output_index as usize, 0);
130 }
131 if let Some(mut queries) = self.static_queries.take() {
132 let n = self.array.len();
133 if queries.is_empty() {
134 return;
135 }
136 let mut offsets = vec![0; n + 2];
137 for &(left, right, _, _) in &queries {
138 if left < right {
139 offsets[left as usize + 1] += 1;
140 offsets[right as usize + 1] += 1;
141 }
142 }
143 for i in 0..=n {
144 offsets[i + 1] += offsets[i];
145 }
146 let mut next = offsets.clone();
147 let mut events = vec![0u32; 2 * queries.len()];
148 for (i, &(left, right, _, output_index)) in queries.iter().enumerate() {
149 if left >= right {
150 callback(output_index as usize, 0);
151 continue;
152 }
153 for (side, position) in [left, right].into_iter().enumerate() {
154 events[next[position as usize]] = (2 * i + side) as u32;
155 next[position as usize] += 1;
156 }
157 }
158 let mut count = vec![0u32; self.values.len()];
159 for position in 0..=n {
160 for &endpoint in &events[offsets[position]..offsets[position + 1]] {
161 let query = (endpoint >> 1) as usize;
162 let frequency = count[queries[query].2 as usize];
163 if endpoint & 1 == 0 {
164 queries[query].0 = frequency;
165 } else {
166 callback(
167 queries[query].3 as usize,
168 (frequency - queries[query].0) as usize,
169 );
170 }
171 }
172 if position < n {
173 count[self.array[position] as usize] += 1;
174 }
175 }
176 return;
177 }
178 let mut processor = RangeFrequencyProcessor::new(self.array.len());
179 for (index, value) in self.array.into_iter().enumerate() {
180 self.events.push((
181 value,
182 RangeFrequencyQuery::Remove {
183 index: index as u32,
184 },
185 ));
186 }
187 let mut offsets = vec![0; self.queried.len() + 1];
188 for &(value, _) in &self.events {
189 offsets[value as usize + 1] += self.queried[value as usize] as usize;
190 }
191 for i in 0..self.queried.len() {
192 offsets[i + 1] += offsets[i];
193 }
194 let mut next = offsets.clone();
195 let mut events = vec![RangeFrequencyQuery::Add { index: 0 }; *offsets.last().unwrap()];
196 for (value, event) in self.events {
197 let value = value as usize;
198 if self.queried[value] != 0 {
199 events[next[value]] = event;
200 next[value] += 1;
201 }
202 }
203 for range in offsets.windows(2) {
204 for &query in &events[range[0]..range[1]] {
205 match query {
206 RangeFrequencyQuery::Add { index } => {
207 processor.add(index);
208 }
209 RangeFrequencyQuery::Remove { index } => {
210 processor.remove(index);
211 }
212 RangeFrequencyQuery::Query {
213 left,
214 right,
215 output_index,
216 } => {
217 callback(output_index as usize, processor.query(left, right));
218 }
219 }
220 }
221 }
222 }
223
224 pub fn execute(self) -> Vec<usize> {
225 let mut results = vec![0; self.output_size];
226 self.execute_with_callback(|i, v| results[i] = v);
227 results
228 }
229}
230
231#[derive(Debug, Clone)]
232struct RangeFrequencyProcessor {
233 bit: BinaryIndexedTree<AdditiveOperation<i32>>,
234 data: Vec<u64>,
235}
236
237impl RangeFrequencyProcessor {
238 fn new(size: usize) -> Self {
239 Self {
240 bit: BinaryIndexedTree::new(size.div_ceil(64)),
241 data: vec![0; size.div_ceil(64)],
242 }
243 }
244
245 fn add(&mut self, index: u32) {
246 let index = index as usize;
247 let (block, bit) = (index / 64, index % 64);
248 assert!(self.data[block] & (1 << bit) == 0);
249 self.data[block] |= 1 << bit;
250 self.bit.update(block, 1);
251 }
252
253 fn remove(&mut self, index: u32) {
254 let index = index as usize;
255 let (i, j) = (index / 64, index % 64);
256 assert!(self.data[i] & (1 << j) != 0);
257 self.data[i] &= !(1 << j);
258 self.bit.update(i, -1);
259 }
260
261 fn query(&self, left: u32, right: u32) -> usize {
262 if left >= right {
263 return 0;
264 }
265 let (left, right) = (left as usize, right as usize - 1);
266 let (li, lj) = (left / 64, left % 64);
267 let (ri, rj) = (right / 64, right % 64);
268 let rj_r = 63 - rj;
269 if li == ri {
270 (self.data[li] << rj_r >> (lj + rj_r)).count_ones() as usize
271 } else {
272 let mut ans = self.bit.fold(li + 1, ri) as usize;
273 ans += (self.data[li] >> lj).count_ones() as usize;
274 ans += (self.data[ri] << rj_r).count_ones() as usize;
275 ans
276 }
277 }
278}
279
280#[cfg(test)]
281mod tests {
282 use super::*;
283 use crate::{
284 rand,
285 tools::{NotEmptySegment as Nes, Xorshift},
286 };
287
288 #[test]
289 fn test_range_frequency() {
290 let mut rng = Xorshift::default();
291 for _ in 0..100 {
292 rand!(rng, n: 1..200, mut a: [0..20; n]);
293 let mut rf = RangeFrequency::new(a.clone());
294 let mut expected = vec![];
295 let dynamic = rng.gen_bool(0.5);
296 for _ in 0..100 {
297 if dynamic && rng.gen_bool(0.5) {
298 rand!(rng, i: 0..n, v: 0..20);
299 rf.set(i, v);
300 a[i] = v;
301 }
302 let (l, r) = rng.random(Nes(n));
303 for v in 0..20 {
304 expected.push(a[l..r].iter().filter(|&&x| x == v).count());
305 rf.query(l, r, v);
306 }
307 }
308 let result = rf.execute();
309 assert_eq!(result, expected);
310 }
311 }
312}