Skip to content

Commit

Permalink
Fuse the iterator in k_smallest
Browse files Browse the repository at this point in the history
  • Loading branch information
Philippe-Cholet authored and phimuemue committed Mar 14, 2024
1 parent 8ed734b commit 1af6c66
Showing 1 changed file with 4 additions and 3 deletions.
7 changes: 4 additions & 3 deletions src/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -2949,7 +2949,7 @@ pub trait Itertools: Iterator {
/// itertools::assert_equal(five_smallest, 0..5);
/// ```
#[cfg(feature = "use_alloc")]
fn k_smallest(mut self, k: usize) -> VecIntoIter<Self::Item>
fn k_smallest(self, k: usize) -> VecIntoIter<Self::Item>
where
Self: Sized,
Self::Item: Ord,
Expand All @@ -2963,9 +2963,10 @@ pub trait Itertools: Iterator {
return Vec::new().into_iter();
}

let mut heap = self.by_ref().take(k).collect::<BinaryHeap<_>>();
let mut iter = self.fuse();
let mut heap: BinaryHeap<_> = iter.by_ref().take(k).collect();

self.for_each(|i| {
iter.for_each(|i| {
debug_assert_eq!(heap.len(), k);
// Equivalent to heap.push(min(i, heap.pop())) but more efficient.
// This should be done with a single `.peek_mut().unwrap()` but
Expand Down

0 comments on commit 1af6c66

Please sign in to comment.