itertools/
kmerge_impl.rs

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

itertools/
kmerge_impl.rs
