1 | use alloc::collections::BinaryHeap; |
2 | use core::cmp::Ord; |
3 | |
4 | pub(crate) fn k_smallest<T: Ord, I: Iterator<Item = T>>(mut iter: I, k: usize) -> BinaryHeap<T> { |
5 | if k == 0 { return BinaryHeap::new(); } |
6 | |
7 | let mut heap = iter.by_ref().take(k).collect::<BinaryHeap<_>>(); |
8 | |
9 | iter.for_each(|i| { |
10 | debug_assert_eq!(heap.len(), k); |
11 | // Equivalent to heap.push(min(i, heap.pop())) but more efficient. |
12 | // This should be done with a single `.peek_mut().unwrap()` but |
13 | // `PeekMut` sifts-down unconditionally on Rust 1.46.0 and prior. |
14 | if *heap.peek().unwrap() > i { |
15 | *heap.peek_mut().unwrap() = i; |
16 | } |
17 | }); |
18 | |
19 | heap |
20 | } |
21 | |