/rust/registry/src/index.crates.io-1949cf8c6b5b557f/itertools-0.11.0/src/k_smallest.rs
Line | Count | Source |
1 | | use alloc::collections::BinaryHeap; |
2 | | use core::cmp::Ord; |
3 | | |
4 | 0 | pub(crate) fn k_smallest<T: Ord, I: Iterator<Item = T>>(mut iter: I, k: usize) -> BinaryHeap<T> { |
5 | 0 | if k == 0 { return BinaryHeap::new(); } |
6 | | |
7 | 0 | let mut heap = iter.by_ref().take(k).collect::<BinaryHeap<_>>(); |
8 | | |
9 | 0 | iter.for_each(|i| { |
10 | 0 | 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 | 0 | if *heap.peek().unwrap() > i { |
15 | 0 | *heap.peek_mut().unwrap() = i; |
16 | 0 | } |
17 | 0 | }); |
18 | | |
19 | 0 | heap |
20 | 0 | } |