1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
use std::cmp::Reverse;
use std::collections::BinaryHeap;
use std::num::NonZeroUsize;
use std::vec::IntoIter as VecIntoIter;
use bytemuck::{TransparentWrapper as _, TransparentWrapperAlloc as _};
use serde::{Deserialize, Serialize};
/// To avoid excessive memory allocation, FixedLengthPriorityQueue
/// imposes a reasonable limit on the allocation size. If the limit
/// is extremely large, we treat it as if no limit was set and
/// delay allocation, assuming that the results will fit within a
/// predefined threshold.
const LARGEST_REASONABLE_ALLOCATION_SIZE: usize = 1_048_576;
/// A container that forgets all but the top N elements
///
/// This is a MinHeap by default - it will keep the largest elements, pop smallest
#[derive(Deserialize, Serialize, Clone, Debug)]
pub struct FixedLengthPriorityQueue<T: Ord> {
heap: BinaryHeap<Reverse<T>>,
length: NonZeroUsize,
}
impl<T: Ord> Default for FixedLengthPriorityQueue<T> {
fn default() -> Self {
Self::new(1)
}
}
impl<T: Ord> FixedLengthPriorityQueue<T> {
/// Creates a new queue with the given length
/// Panics if length is 0
pub fn new(length: usize) -> Self {
let heap = BinaryHeap::with_capacity(
length
.saturating_add(1)
.min(LARGEST_REASONABLE_ALLOCATION_SIZE),
);
let length = NonZeroUsize::new(length).expect("length must be greater than zero");
FixedLengthPriorityQueue::<T> { heap, length }
}
/// Pushes a value into the priority queue.
///
/// If the queue if full, replaces the smallest value and returns it.
pub fn push(&mut self, value: T) -> Option<T> {
if !self.is_full() {
self.heap.push(Reverse(value));
return None;
}
let mut x = self.heap.peek_mut().unwrap();
let mut value = Reverse(value);
if x.0 < value.0 {
std::mem::swap(&mut *x, &mut value);
}
Some(value.0)
}
/// Consumes the [`FixedLengthPriorityQueue`] and returns a vector
/// in sorted (descending) order.
pub fn into_sorted_vec(self) -> Vec<T> {
Reverse::peel_vec(self.heap.into_sorted_vec())
}
/// Returns an iterator over the elements in the queue, in arbitrary order.
pub fn iter_unsorted(&self) -> std::slice::Iter<'_, T> {
Reverse::peel_slice(self.heap.as_slice()).iter()
}
/// Returns an iterator over the elements in the queue
/// in sorted (descending) order.
pub fn into_iter_sorted(self) -> VecIntoIter<T> {
self.into_sorted_vec().into_iter()
}
/// Returns the smallest element of the queue,
/// if there is any.
pub fn top(&self) -> Option<&T> {
self.heap.peek().map(|x| &x.0)
}
/// Returns actual length of the queue
pub fn len(&self) -> usize {
self.heap.len()
}
/// Checks if the queue is empty
pub fn is_empty(&self) -> bool {
self.heap.is_empty()
}
/// Checks if the queue is full
pub fn is_full(&self) -> bool {
self.heap.len() >= self.length.into()
}
pub fn retain<F>(&mut self, mut f: F)
where
F: FnMut(&T) -> bool,
{
self.heap.retain(|x| f(&x.0));
}
}