itertools/
k_smallest.rs

1use alloc::vec::Vec;
2use core::cmp::Ordering;
3
4/// Consumes a given iterator, returning the minimum elements in **ascending** order.
5pub(crate) fn k_smallest_general<I, F>(iter: I, k: usize, mut comparator: F) -> Vec<I::Item>
6where
7    I: Iterator,
8    F: FnMut(&I::Item, &I::Item) -> Ordering,
9{
10    /// Sift the element currently at `origin` away from the root until it is properly ordered.
11    ///
12    /// This will leave **larger** elements closer to the root of the heap.
13    fn sift_down<T, F>(heap: &mut [T], is_less_than: &mut F, mut origin: usize)
14    where
15        F: FnMut(&T, &T) -> bool,
16    {
17        #[inline]
18        fn children_of(n: usize) -> (usize, usize) {
19            (2 * n + 1, 2 * n + 2)
20        }
21
22        while origin < heap.len() {
23            let (left_idx, right_idx) = children_of(origin);
24            if left_idx >= heap.len() {
25                return;
26            }
27
28            let replacement_idx =
29                if right_idx < heap.len() && is_less_than(&heap[left_idx], &heap[right_idx]) {
30                    right_idx
31                } else {
32                    left_idx
33                };
34
35            if is_less_than(&heap[origin], &heap[replacement_idx]) {
36                heap.swap(origin, replacement_idx);
37                origin = replacement_idx;
38            } else {
39                return;
40            }
41        }
42    }
43
44    if k == 0 {
45        iter.last();
46        return Vec::new();
47    }
48    if k == 1 {
49        return iter.min_by(comparator).into_iter().collect();
50    }
51    let mut iter = iter.fuse();
52    let mut storage: Vec<I::Item> = iter.by_ref().take(k).collect();
53
54    let mut is_less_than = move |a: &_, b: &_| comparator(a, b) == Ordering::Less;
55
56    // Rearrange the storage into a valid heap by reordering from the second-bottom-most layer up to the root.
57    // Slightly faster than ordering on each insert, but only by a factor of lg(k).
58    // The resulting heap has the **largest** item on top.
59    for i in (0..=(storage.len() / 2)).rev() {
60        sift_down(&mut storage, &mut is_less_than, i);
61    }
62
63    iter.for_each(|val| {
64        debug_assert_eq!(storage.len(), k);
65        if is_less_than(&val, &storage[0]) {
66            // Treating this as an push-and-pop saves having to write a sift-up implementation.
67            // https://en.wikipedia.org/wiki/Binary_heap#Insert_then_extract
68            storage[0] = val;
69            // We retain the smallest items we've seen so far, but ordered largest first so we can drop the largest efficiently.
70            sift_down(&mut storage, &mut is_less_than, 0);
71        }
72    });
73
74    // Ultimately the items need to be in least-first, strict order, but the heap is currently largest-first.
75    // To achieve this, repeatedly,
76    // 1) "pop" the largest item off the heap into the tail slot of the underlying storage,
77    // 2) shrink the logical size of the heap by 1,
78    // 3) restore the heap property over the remaining items.
79    let mut heap = &mut storage[..];
80    while heap.len() > 1 {
81        let last_idx = heap.len() - 1;
82        heap.swap(0, last_idx);
83        // Sifting over a truncated slice means that the sifting will not disturb already popped elements.
84        heap = &mut heap[..last_idx];
85        sift_down(heap, &mut is_less_than, 0);
86    }
87
88    storage
89}
90
91pub(crate) fn k_smallest_relaxed_general<I, F>(iter: I, k: usize, mut comparator: F) -> Vec<I::Item>
92where
93    I: Iterator,
94    F: FnMut(&I::Item, &I::Item) -> Ordering,
95{
96    if k == 0 {
97        iter.last();
98        return Vec::new();
99    }
100
101    let mut iter = iter.fuse();
102    let mut buf = iter.by_ref().take(2 * k).collect::<Vec<_>>();
103
104    if buf.len() < k {
105        buf.sort_unstable_by(&mut comparator);
106        return buf;
107    }
108
109    buf.select_nth_unstable_by(k - 1, &mut comparator);
110    buf.truncate(k);
111
112    iter.for_each(|val| {
113        if comparator(&val, &buf[k - 1]) != Ordering::Less {
114            return;
115        }
116
117        assert_ne!(buf.len(), buf.capacity());
118        buf.push(val);
119
120        if buf.len() == 2 * k {
121            buf.select_nth_unstable_by(k - 1, &mut comparator);
122            buf.truncate(k);
123        }
124    });
125
126    buf.sort_unstable_by(&mut comparator);
127    buf.truncate(k);
128    buf
129}
130
131#[inline]
132pub(crate) fn key_to_cmp<T, K, F>(mut key: F) -> impl FnMut(&T, &T) -> Ordering
133where
134    F: FnMut(&T) -> K,
135    K: Ord,
136{
137    move |a, b| key(a).cmp(&key(b))
138}
itertools/k_smallest.rs

itertools/
k_smallest.rs
