Add Itertool::k_smallest method, in std context

Author: nicoonoclaste

src/k_smallest.rs

@@ -0,0 +1,16 @@
+use std::collections::BinaryHeap;
+use std::cmp::Ord;
+
+pub(crate) fn k_smallest<T: Ord, I: Iterator<Item = T>>(mut iter: I, k: usize) -> BinaryHeap<T> {
+    if k == 0 { return BinaryHeap::new(); }
+
+    let mut heap = iter.by_ref().take(k).collect::<BinaryHeap<_>>();
+
+    for i in iter {
+        // Guaranteed not-None, since we keep exactly k>0 elements in the heap.
+        let mut lorgest = heap.peek_mut().unwrap();
+        if *lorgest > i { *lorgest = i; }
+    }
+
+    heap
+}

src/lib.rs

@@ -194,6 +194,8 @@ mod group_map;
 mod groupbylazy;
 mod intersperse;
 #[cfg(feature = "use_alloc")]
+mod k_smallest;
+#[cfg(feature = "use_alloc")]
 mod kmerge_impl;
 #[cfg(feature = "use_alloc")]
 mod lazy_buffer;
@@ -2284,6 +2286,39 @@ pub trait Itertools : Iterator {
         v.into_iter()
     }
 
+    /// Sort the k smallest elements into a new iterator, in ascending order.
+    ///
+    /// **Note:** This consumes the entire iterator, and returns the result
+    /// as a new iterator that owns its elements.
+    ///
+    /// This is guaranteed to use `k * sizeof(Self::Item) + O(1)` memory
+    /// and `O(n log k)` time, with `n` the number of elements in the input.
+    ///
+    /// The sorted iterator, if directly collected to a `Vec`, is converted
+    /// without any extra copying or allocation cost.
+    ///
+    /// ```
+    /// use itertools::Itertools;
+    ///
+    /// // A random permutation of 0..15
+    /// let numbers = vec![6, 9, 1, 14, 0, 4, 8, 7, 11, 2, 10, 3, 13, 12, 5];
+    ///
+    /// let five_smallest = numbers
+    ///     .into_iter()
+    ///     .k_smallest(5);
+    ///
+    /// itertools::assert_equal(five_smallest, 0..5);
+    /// ```
+    #[cfg(feature = "use_alloc")]
+    fn k_smallest(self, k: usize) -> VecIntoIter<Self::Item>
+        where Self: Sized,
+              Self::Item: Ord
+    {
+        crate::k_smallest::k_smallest(self, k)
+            .into_sorted_vec()
+            .into_iter()
+    }
+
     /// Collect all iterator elements into one of two
     /// partitions. Unlike `Iterator::partition`, each partition may
     /// have a distinct type.