itertools/
kmerge_impl.rs

1use crate::size_hint;
2
3use alloc::vec::Vec;
4use std::fmt;
5use std::iter::FusedIterator;
6use std::mem::replace;
7
8/// Head element and Tail iterator pair
9///
10/// `PartialEq`, `Eq`, `PartialOrd` and `Ord` are implemented by comparing sequences based on
11/// first items (which are guaranteed to exist).
12///
13/// The meanings of `PartialOrd` and `Ord` are reversed so as to turn the heap used in
14/// `KMerge` into a min-heap.
15#[derive(Debug)]
16struct HeadTail<I>
17where
18    I: Iterator,
19{
20    head: I::Item,
21    tail: I,
22}
23
24impl<I> HeadTail<I>
25where
26    I: Iterator,
27{
28    /// Constructs a `HeadTail` from an `Iterator`. Returns `None` if the `Iterator` is empty.
29    fn new(mut it: I) -> Option<Self> {
30        let head = it.next();
31        head.map(|h| Self { head: h, tail: it })
32    }
33
34    /// Get the next element and update `head`, returning the old head in `Some`.
35    ///
36    /// Returns `None` when the tail is exhausted (only `head` then remains).
37    fn next(&mut self) -> Option<I::Item> {
38        if let Some(next) = self.tail.next() {
39            Some(replace(&mut self.head, next))
40        } else {
41            None
42        }
43    }
44
45    /// Hints at the size of the sequence, same as the `Iterator` method.
46    fn size_hint(&self) -> (usize, Option<usize>) {
47        size_hint::add_scalar(self.tail.size_hint(), 1)
48    }
49}
50
51impl<I> Clone for HeadTail<I>
52where
53    I: Iterator + Clone,
54    I::Item: Clone,
55{
56    clone_fields!(head, tail);
57}
58
59/// Make `data` a heap (min-heap w.r.t the sorting).
60fn heapify<T, S>(data: &mut [T], mut less_than: S)
61where
62    S: FnMut(&T, &T) -> bool,
63{
64    for i in (0..data.len() / 2).rev() {
65        sift_down(data, i, &mut less_than);
66    }
67}
68
69/// Sift down element at `index` (`heap` is a min-heap wrt the ordering)
70fn sift_down<T, S>(heap: &mut [T], index: usize, mut less_than: S)
71where
72    S: FnMut(&T, &T) -> bool,
73{
74    debug_assert!(index <= heap.len());
75    let mut pos = index;
76    let mut child = 2 * pos + 1;
77    // Require the right child to be present
78    // This allows to find the index of the smallest child without a branch
79    // that wouldn't be predicted if present
80    while child + 1 < heap.len() {
81        // pick the smaller of the two children
82        // use arithmetic to avoid an unpredictable branch
83        child += less_than(&heap[child + 1], &heap[child]) as usize;
84
85        // sift down is done if we are already in order
86        if !less_than(&heap[child], &heap[pos]) {
87            return;
88        }
89        heap.swap(pos, child);
90        pos = child;
91        child = 2 * pos + 1;
92    }
93    // Check if the last (left) child was an only child
94    // if it is then it has to be compared with the parent
95    if child + 1 == heap.len() && less_than(&heap[child], &heap[pos]) {
96        heap.swap(pos, child);
97    }
98}
99
100/// An iterator adaptor that merges an abitrary number of base iterators in ascending order.
101/// If all base iterators are sorted (ascending), the result is sorted.
102///
103/// Iterator element type is `I::Item`.
104///
105/// See [`.kmerge()`](crate::Itertools::kmerge) for more information.
106pub type KMerge<I> = KMergeBy<I, KMergeByLt>;
107
108pub trait KMergePredicate<T> {
109    fn kmerge_pred(&mut self, a: &T, b: &T) -> bool;
110}
111
112#[derive(Clone, Debug)]
113pub struct KMergeByLt;
114
115impl<T: PartialOrd> KMergePredicate<T> for KMergeByLt {
116    fn kmerge_pred(&mut self, a: &T, b: &T) -> bool {
117        a < b
118    }
119}
120
121impl<T, F: FnMut(&T, &T) -> bool> KMergePredicate<T> for F {
122    fn kmerge_pred(&mut self, a: &T, b: &T) -> bool {
123        self(a, b)
124    }
125}
126
127/// Create an iterator that merges elements of the contained iterators using
128/// the ordering function.
129///
130/// [`IntoIterator`] enabled version of [`Itertools::kmerge`](crate::Itertools::kmerge).
131///
132/// ```
133/// use itertools::kmerge;
134///
135/// for elt in kmerge(vec![vec![0, 2, 4], vec![1, 3, 5], vec![6, 7]]) {
136///     /* loop body */
137///     # let _ = elt;
138/// }
139/// ```
140pub fn kmerge<I>(iterable: I) -> KMerge<<I::Item as IntoIterator>::IntoIter>
141where
142    I: IntoIterator,
143    I::Item: IntoIterator,
144    <<I as IntoIterator>::Item as IntoIterator>::Item: PartialOrd,
145{
146    kmerge_by(iterable, KMergeByLt)
147}
148
149/// An iterator adaptor that merges an abitrary number of base iterators
150/// according to an ordering function.
151///
152/// Iterator element type is `I::Item`.
153///
154/// See [`.kmerge_by()`](crate::Itertools::kmerge_by) for more
155/// information.
156#[must_use = "this iterator adaptor is not lazy but does nearly nothing unless consumed"]
157pub struct KMergeBy<I, F>
158where
159    I: Iterator,
160{
161    heap: Vec<HeadTail<I>>,
162    less_than: F,
163}
164
165impl<I, F> fmt::Debug for KMergeBy<I, F>
166where
167    I: Iterator + fmt::Debug,
168    I::Item: fmt::Debug,
169{
170    debug_fmt_fields!(KMergeBy, heap);
171}
172
173/// Create an iterator that merges elements of the contained iterators.
174///
175/// [`IntoIterator`] enabled version of [`Itertools::kmerge_by`](crate::Itertools::kmerge_by).
176pub fn kmerge_by<I, F>(
177    iterable: I,
178    mut less_than: F,
179) -> KMergeBy<<I::Item as IntoIterator>::IntoIter, F>
180where
181    I: IntoIterator,
182    I::Item: IntoIterator,
183    F: KMergePredicate<<<I as IntoIterator>::Item as IntoIterator>::Item>,
184{
185    let iter = iterable.into_iter();
186    let (lower, _) = iter.size_hint();
187    let mut heap: Vec<_> = Vec::with_capacity(lower);
188    heap.extend(iter.filter_map(|it| HeadTail::new(it.into_iter())));
189    heapify(&mut heap, |a, b| less_than.kmerge_pred(&a.head, &b.head));
190    KMergeBy { heap, less_than }
191}
192
193impl<I, F> Clone for KMergeBy<I, F>
194where
195    I: Iterator + Clone,
196    I::Item: Clone,
197    F: Clone,
198{
199    clone_fields!(heap, less_than);
200}
201
202impl<I, F> Iterator for KMergeBy<I, F>
203where
204    I: Iterator,
205    F: KMergePredicate<I::Item>,
206{
207    type Item = I::Item;
208
209    fn next(&mut self) -> Option<Self::Item> {
210        if self.heap.is_empty() {
211            return None;
212        }
213        let result = if let Some(next) = self.heap[0].next() {
214            next
215        } else {
216            self.heap.swap_remove(0).head
217        };
218        let less_than = &mut self.less_than;
219        sift_down(&mut self.heap, 0, |a, b| {
220            less_than.kmerge_pred(&a.head, &b.head)
221        });
222        Some(result)
223    }
224
225    fn size_hint(&self) -> (usize, Option<usize>) {
226        self.heap
227            .iter()
228            .map(|i| i.size_hint())
229            .reduce(size_hint::add)
230            .unwrap_or((0, Some(0)))
231    }
232}
233
234impl<I, F> FusedIterator for KMergeBy<I, F>
235where
236    I: Iterator,
237    F: KMergePredicate<I::Item>,
238{
239}
itertools/kmerge_impl.rs

itertools/
kmerge_impl.rs
