competitive/data_structure/
bucket_queue.rs1#[derive(Clone, Debug)]
2struct BucketQueue8 {
3 counts: [u32; 1 << 8],
4 occupied: [u64; 1 << 2],
5 summary: u8,
6 maximum: u8,
7 len: usize,
8}
9
10impl BucketQueue8 {
11 fn new() -> Self {
12 Self {
13 counts: [0; 1 << 8],
14 occupied: [0; 1 << 2],
15 summary: 0,
16 maximum: 0,
17 len: 0,
18 }
19 }
20
21 #[inline]
22 fn push(&mut self, value: u8) {
23 assert!(self.len < u32::MAX as usize);
24 let value = value as usize;
25 if self.len == 0 || value > self.maximum as usize {
26 self.maximum = value as u8;
27 }
28 if self.counts[value] == 0 {
29 self.occupied[value / 64] |= 1 << (value % 64);
30 self.summary |= 1 << (value / 64);
31 }
32 self.counts[value] += 1;
33 self.len += 1;
34 }
35
36 fn from_values(values: impl IntoIterator<Item = u8>, len: usize) -> Self {
37 assert!(len <= u32::MAX as usize);
38 let mut result = Self::new();
39 result.len = len;
40 for value in values {
41 result.counts[value as usize] += 1;
42 }
43 for (value, &count) in result.counts.iter().enumerate() {
44 if count != 0 {
45 result.occupied[value / 64] |= 1 << (value % 64);
46 }
47 }
48 for (word, &occupied) in result.occupied.iter().enumerate() {
49 if occupied != 0 {
50 result.summary |= 1 << word;
51 }
52 }
53 if len != 0 {
54 let word = (u8::BITS - 1 - result.summary.leading_zeros()) as usize;
55 result.maximum =
56 (word * 64 + 63 - result.occupied[word].leading_zeros() as usize) as u8;
57 }
58 result
59 }
60
61 #[inline]
62 fn pop(&mut self) -> Option<u8> {
63 if self.len == 0 {
64 return None;
65 }
66 let value = self.maximum as usize;
67 self.counts[value] -= 1;
68 self.len -= 1;
69 if self.counts[value] == 0 {
70 let word = value / 64;
71 self.occupied[word] &= !(1 << (value % 64));
72 if self.occupied[word] == 0 {
73 self.summary &= !(1 << word);
74 }
75 if self.len != 0 {
76 let word = (u8::BITS - 1 - self.summary.leading_zeros()) as usize;
77 self.maximum =
78 (word * 64 + 63 - self.occupied[word].leading_zeros() as usize) as u8;
79 }
80 }
81 Some(value as u8)
82 }
83
84 #[inline]
85 fn replace(&mut self, value: u8) -> Option<u8> {
86 if self.len == 0 {
87 self.push(value);
88 return None;
89 }
90 let result = self.maximum;
91 if value == result {
92 return Some(result);
93 }
94
95 let old = result as usize;
96 self.counts[old] -= 1;
97 if self.counts[old] == 0 {
98 let word = old / 64;
99 self.occupied[word] &= !(1 << (old % 64));
100 if self.occupied[word] == 0 {
101 self.summary &= !(1 << word);
102 }
103 }
104
105 let new = value as usize;
106 if self.counts[new] == 0 {
107 self.occupied[new / 64] |= 1 << (new % 64);
108 self.summary |= 1 << (new / 64);
109 }
110 self.counts[new] += 1;
111
112 if value > result || self.counts[old] == 0 {
113 let word = (u8::BITS - 1 - self.summary.leading_zeros()) as usize;
114 self.maximum = (word * 64 + 63 - self.occupied[word].leading_zeros() as usize) as u8;
115 }
116 Some(result)
117 }
118
119 fn clear(&mut self) {
120 self.counts.fill(0);
121 self.occupied.fill(0);
122 self.summary = 0;
123 self.maximum = 0;
124 self.len = 0;
125 }
126}
127
128#[derive(Clone, Debug)]
129struct BucketQueue16 {
130 counts: Vec<u32>,
131 occupied: Vec<u64>,
132 summary: [u64; 1 << 4],
133 top: u16,
134 maximum: u16,
135 len: usize,
136}
137
138impl BucketQueue16 {
139 fn new() -> Self {
140 Self {
141 counts: vec![0; 1 << 16],
142 occupied: vec![0; 1 << 10],
143 summary: [0; 1 << 4],
144 top: 0,
145 maximum: 0,
146 len: 0,
147 }
148 }
149
150 #[inline]
151 fn push(&mut self, value: u16) {
152 assert!(self.len < u32::MAX as usize);
153 let value = value as usize;
154 if self.len == 0 || value > self.maximum as usize {
155 self.maximum = value as u16;
156 }
157 if self.counts[value] == 0 {
158 let word = value / 64;
159 self.occupied[word] |= 1 << (value % 64);
160 self.summary[word / 64] |= 1 << (word % 64);
161 self.top |= 1 << (word / 64);
162 }
163 self.counts[value] += 1;
164 self.len += 1;
165 }
166
167 fn from_values(values: impl IntoIterator<Item = u16>, len: usize) -> Self {
168 assert!(len <= u32::MAX as usize);
169 let mut result = Self::new();
170 result.len = len;
171 for value in values {
172 result.counts[value as usize] += 1;
173 }
174 for (value, &count) in result.counts.iter().enumerate() {
175 if count != 0 {
176 result.occupied[value / 64] |= 1 << (value % 64);
177 }
178 }
179 for (word, &occupied) in result.occupied.iter().enumerate() {
180 if occupied != 0 {
181 result.summary[word / 64] |= 1 << (word % 64);
182 }
183 }
184 for (word, &summary) in result.summary.iter().enumerate() {
185 if summary != 0 {
186 result.top |= 1 << word;
187 }
188 }
189 if len != 0 {
190 let summary = (u16::BITS - 1 - result.top.leading_zeros()) as usize;
191 let word = summary * 64 + 63 - result.summary[summary].leading_zeros() as usize;
192 result.maximum =
193 (word * 64 + 63 - result.occupied[word].leading_zeros() as usize) as u16;
194 }
195 result
196 }
197
198 #[inline]
199 fn pop(&mut self) -> Option<u16> {
200 if self.len == 0 {
201 return None;
202 }
203 let value = self.maximum as usize;
204 self.counts[value] -= 1;
205 self.len -= 1;
206 if self.counts[value] == 0 {
207 let word = value / 64;
208 let summary = word / 64;
209 self.occupied[word] &= !(1 << (value % 64));
210 if self.occupied[word] == 0 {
211 self.summary[summary] &= !(1 << (word % 64));
212 if self.summary[summary] == 0 {
213 self.top &= !(1 << summary);
214 }
215 }
216 if self.len != 0 {
217 let summary = (u16::BITS - 1 - self.top.leading_zeros()) as usize;
218 let word = summary * 64 + 63 - self.summary[summary].leading_zeros() as usize;
219 self.maximum =
220 (word * 64 + 63 - self.occupied[word].leading_zeros() as usize) as u16;
221 }
222 }
223 Some(value as u16)
224 }
225
226 #[inline]
227 fn replace(&mut self, value: u16) -> Option<u16> {
228 if self.len == 0 {
229 self.push(value);
230 return None;
231 }
232 let result = self.maximum;
233 if value == result {
234 return Some(result);
235 }
236
237 let old = result as usize;
238 self.counts[old] -= 1;
239 if self.counts[old] == 0 {
240 let word = old / 64;
241 let summary = word / 64;
242 self.occupied[word] &= !(1 << (old % 64));
243 if self.occupied[word] == 0 {
244 self.summary[summary] &= !(1 << (word % 64));
245 if self.summary[summary] == 0 {
246 self.top &= !(1 << summary);
247 }
248 }
249 }
250
251 let new = value as usize;
252 if self.counts[new] == 0 {
253 let word = new / 64;
254 self.occupied[word] |= 1 << (new % 64);
255 self.summary[word / 64] |= 1 << (word % 64);
256 self.top |= 1 << (word / 64);
257 }
258 self.counts[new] += 1;
259
260 if value > result || self.counts[old] == 0 {
261 let summary = (u16::BITS - 1 - self.top.leading_zeros()) as usize;
262 let word = summary * 64 + 63 - self.summary[summary].leading_zeros() as usize;
263 self.maximum = (word * 64 + 63 - self.occupied[word].leading_zeros() as usize) as u16;
264 }
265 Some(result)
266 }
267
268 fn clear(&mut self) {
269 self.counts[..=self.maximum as usize].fill(0);
270 self.occupied.fill(0);
271 self.summary.fill(0);
272 self.top = 0;
273 self.maximum = 0;
274 self.len = 0;
275 }
276}
277
278macro_rules! define_bucket_queue {
279 ($name:ident, $doc:literal, $value:ty, $repr:ty, $queue:ty, $sign:expr, $bulk_threshold:expr) => {
280 #[doc = $doc]
281 #[derive(Clone, Debug)]
282 pub struct $name {
283 queue: $queue,
284 }
285
286 impl $name {
287 pub fn new() -> Self {
288 Self {
289 queue: <$queue>::new(),
290 }
291 }
292
293 #[inline]
294 pub fn len(&self) -> usize {
295 self.queue.len
296 }
297
298 #[inline]
299 pub fn is_empty(&self) -> bool {
300 self.queue.len == 0
301 }
302
303 #[inline]
304 pub fn peek(&self) -> Option<$value> {
305 (self.queue.len != 0).then_some((self.queue.maximum ^ $sign) as $value)
306 }
307
308 #[inline]
312 pub fn push(&mut self, value: $value) {
313 self.queue.push((value as $repr) ^ $sign);
314 }
315
316 #[inline]
317 pub fn pop(&mut self) -> Option<$value> {
318 self.queue.pop().map(|value| ((value ^ $sign) as $value))
319 }
320
321 #[inline]
323 pub fn replace(&mut self, value: $value) -> Option<$value> {
324 self.queue
325 .replace((value as $repr) ^ $sign)
326 .map(|value| (value ^ $sign) as $value)
327 }
328
329 pub fn clear(&mut self) {
330 self.queue.clear();
331 }
332 }
333
334 impl Default for $name {
335 fn default() -> Self {
336 Self::new()
337 }
338 }
339
340 impl From<Vec<$value>> for $name {
341 fn from(values: Vec<$value>) -> Self {
342 if values.len() >= $bulk_threshold {
343 let len = values.len();
344 Self {
345 queue: <$queue>::from_values(
346 values.into_iter().map(|value| (value as $repr) ^ $sign),
347 len,
348 ),
349 }
350 } else {
351 let mut queue = Self::new();
352 queue.extend(values);
353 queue
354 }
355 }
356 }
357
358 impl Extend<$value> for $name {
359 fn extend<I>(&mut self, iter: I)
360 where
361 I: IntoIterator<Item = $value>,
362 {
363 for value in iter {
364 self.push(value);
365 }
366 }
367 }
368
369 impl FromIterator<$value> for $name {
370 fn from_iter<I>(iter: I) -> Self
371 where
372 I: IntoIterator<Item = $value>,
373 {
374 Self::from(Vec::from_iter(iter))
375 }
376 }
377 };
378}
379
380define_bucket_queue!(
381 BucketQueueU8,
382 "A fixed 8-bit-universe max-priority queue. `BinaryHeap::peek_mut` can be faster for replacements in tiny queues.",
383 u8,
384 u8,
385 BucketQueue8,
386 0,
387 1 << 12
388);
389define_bucket_queue!(
390 BucketQueueI8,
391 "A fixed 8-bit-universe max-priority queue. `BinaryHeap::peek_mut` can be faster for replacements in tiny queues.",
392 i8,
393 u8,
394 BucketQueue8,
395 1 << 7,
396 1 << 12
397);
398define_bucket_queue!(
399 BucketQueueU16,
400 "A fixed 16-bit-universe max-priority queue that allocates about 264 KiB when empty. `BinaryHeap` can be faster for small queues.",
401 u16,
402 u16,
403 BucketQueue16,
404 0,
405 1 << 16
406);
407define_bucket_queue!(
408 BucketQueueI16,
409 "A fixed 16-bit-universe max-priority queue that allocates about 264 KiB when empty. `BinaryHeap` can be faster for small queues.",
410 i16,
411 u16,
412 BucketQueue16,
413 1 << 15,
414 1 << 16
415);
416
417#[cfg(test)]
418mod tests {
419 use super::*;
420 use crate::tools::Xorshift;
421 use crate::tools::testutil::{exhaustive_sequences, integer_boundary_values};
422 use std::collections::BinaryHeap;
423
424 #[test]
425 fn test_bucket_queue() {
426 macro_rules! test_queue {
427 ($queue:ty, $value:ty) => {{
428 let mut rng = Xorshift::default();
429 let values: Vec<$value> = (0..10_000).map(|_| rng.rand64() as $value).collect();
430 let mut actual: $queue = values.clone().into();
431 let mut expected = BinaryHeap::from(values);
432 for _ in 0..20_000 {
433 match rng.rand(4) {
434 0 => {
435 let value = rng.rand64() as $value;
436 actual.push(value);
437 expected.push(value);
438 }
439 1 => assert_eq!(actual.pop(), expected.pop()),
440 _ => {
441 let value = rng.rand64() as $value;
442 let old = expected.pop();
443 expected.push(value);
444 assert_eq!(actual.replace(value), old);
445 }
446 }
447 assert_eq!(actual.peek(), expected.peek().copied());
448 assert_eq!(actual.len(), expected.len());
449 assert_eq!(actual.is_empty(), expected.is_empty());
450 }
451 while let Some(value) = expected.pop() {
452 assert_eq!(actual.pop(), Some(value));
453 }
454 assert_eq!(actual.pop(), None);
455 let values = integer_boundary_values!($value);
456 actual.extend(values.iter().copied());
457 expected.extend(values);
458 while let Some(value) = expected.pop() {
459 assert_eq!(actual.pop(), Some(value));
460 }
461 for values in exhaustive_sequences([<$value>::MIN, 0, <$value>::MAX], 0..=6)
462 .chain([integer_boundary_values!($value)])
463 {
464 let mut actual: $queue = values.clone().into();
465 let mut cleared = actual.clone();
466 cleared.clear();
467 assert_eq!(cleared.len(), 0);
468 assert!(cleared.is_empty());
469 assert_eq!(cleared.peek(), None);
470 assert_eq!(cleared.pop(), None);
471 cleared.extend(values.iter().copied());
472 let mut expected = BinaryHeap::from(values);
473 while let Some(value) = expected.pop() {
474 assert_eq!(actual.pop(), Some(value));
475 assert_eq!(cleared.pop(), Some(value));
476 }
477 assert_eq!(actual.pop(), None);
478 assert_eq!(cleared.pop(), None);
479 }
480 }};
481 }
482
483 test_queue!(BucketQueueU8, u8);
484 test_queue!(BucketQueueI8, i8);
485 test_queue!(BucketQueueU16, u16);
486 test_queue!(BucketQueueI16, i16);
487
488 let values: Vec<_> = (0..1 << 16).map(|value| value as u16).collect();
489 let mut actual = BucketQueueU16::from(values.clone());
490 let mut expected = BinaryHeap::from(values);
491 while !expected.is_empty() {
492 assert_eq!(actual.pop(), expected.pop());
493 }
494 }
495}