Skip to main content

kaspa_utils/
binary_heap.rs

1use alloc::collections::BinaryHeap;
2use core::cmp::Reverse;
3pub trait BinaryHeapExtensions<T> {
4    fn into_sorted_iter(self) -> BinaryHeapIntoSortedIter<T>;
5}
6
7pub struct BinaryHeapIntoSortedIter<T> {
8    binary_heap: BinaryHeap<T>,
9}
10
11impl<T: Ord> Iterator for BinaryHeapIntoSortedIter<T> {
12    type Item = T;
13
14    fn next(&mut self) -> Option<Self::Item> {
15        self.binary_heap.pop()
16    }
17}
18
19impl<T> BinaryHeapExtensions<T> for BinaryHeap<T> {
20    fn into_sorted_iter(self) -> BinaryHeapIntoSortedIter<T> {
21        BinaryHeapIntoSortedIter { binary_heap: self }
22    }
23}
24
25/// Maintains the top `K` elements seen so far (by `Ord`), using a bounded min-heap.
26///
27/// Internally keeps a `BinaryHeap<Reverse<T>>` so `peek()` returns the current minimum
28/// among the kept items.
29#[derive(Default, Clone)]
30pub struct TopK<T: Ord, const K: usize> {
31    heap: BinaryHeap<Reverse<T>>,
32}
33
34impl<T: Ord, const K: usize> TopK<T, K> {
35    pub fn new() -> Self {
36        Self { heap: BinaryHeap::new() }
37    }
38
39    /// Pushes an item, keeping only the top `K` items seen so far.
40    pub fn push(&mut self, item: T) {
41        if self.heap.len() < K {
42            self.heap.push(Reverse(item));
43        } else if let Some(Reverse(min)) = self.heap.peek()
44            && item > *min
45        {
46            self.heap.pop();
47            self.heap.push(Reverse(item));
48        }
49    }
50
51    /// Consumes `self` and returns an iterator yielding items in ascending order.
52    ///
53    /// This uses `BinaryHeapExtensions::into_sorted_iter`, which pops from the underlying heap.
54    pub fn into_sorted_iter_ascending(self) -> impl Iterator<Item = T> {
55        self.heap.into_sorted_iter().map(|Reverse(x)| x)
56    }
57}