use crate::size_hint; use crate::Itertools; use alloc::vec::Vec; use std::fmt; use std::iter::FusedIterator; use std::mem::replace; /// Head element and Tail iterator pair /// /// `PartialEq`, `Eq`, `PartialOrd` and `Ord` are implemented by comparing sequences based on /// first items (which are guaranteed to exist). /// /// The meanings of `PartialOrd` and `Ord` are reversed so as to turn the heap used in /// `KMerge` into a min-heap. #[derive(Debug)] struct HeadTail where I: Iterator, { head: I::Item, tail: I, } impl HeadTail where I: Iterator, { /// Constructs a `HeadTail` from an `Iterator`. Returns `None` if the `Iterator` is empty. fn new(mut it: I) -> Option> { let head = it.next(); head.map(|h| HeadTail { head: h, tail: it }) } /// Get the next element and update `head`, returning the old head in `Some`. /// /// Returns `None` when the tail is exhausted (only `head` then remains). fn next(&mut self) -> Option { if let Some(next) = self.tail.next() { Some(replace(&mut self.head, next)) } else { None } } /// Hints at the size of the sequence, same as the `Iterator` method. fn size_hint(&self) -> (usize, Option) { size_hint::add_scalar(self.tail.size_hint(), 1) } } impl Clone for HeadTail where I: Iterator + Clone, I::Item: Clone, { clone_fields!(head, tail); } /// Make `data` a heap (min-heap w.r.t the sorting). fn heapify(data: &mut [T], mut less_than: S) where S: FnMut(&T, &T) -> bool, { for i in (0..data.len() / 2).rev() { sift_down(data, i, &mut less_than); } } /// Sift down element at `index` (`heap` is a min-heap wrt the ordering) fn sift_down(heap: &mut [T], index: usize, mut less_than: S) where S: FnMut(&T, &T) -> bool, { debug_assert!(index <= heap.len()); let mut pos = index; let mut child = 2 * pos + 1; // Require the right child to be present // This allows to find the index of the smallest child without a branch // that wouldn't be predicted if present while child + 1 < heap.len() { // pick the smaller of the two children // use arithmetic to avoid an unpredictable branch child += less_than(&heap[child + 1], &heap[child]) as usize; // sift down is done if we are already in order if !less_than(&heap[child], &heap[pos]) { return; } heap.swap(pos, child); pos = child; child = 2 * pos + 1; } // Check if the last (left) child was an only child // if it is then it has to be compared with the parent if child + 1 == heap.len() && less_than(&heap[child], &heap[pos]) { heap.swap(pos, child); } } /// An iterator adaptor that merges an abitrary number of base iterators in ascending order. /// If all base iterators are sorted (ascending), the result is sorted. /// /// Iterator element type is `I::Item`. /// /// See [`.kmerge()`](crate::Itertools::kmerge) for more information. pub type KMerge = KMergeBy; pub trait KMergePredicate { fn kmerge_pred(&mut self, a: &T, b: &T) -> bool; } #[derive(Clone, Debug)] pub struct KMergeByLt; impl KMergePredicate for KMergeByLt { fn kmerge_pred(&mut self, a: &T, b: &T) -> bool { a < b } } impl bool> KMergePredicate for F { fn kmerge_pred(&mut self, a: &T, b: &T) -> bool { self(a, b) } } /// Create an iterator that merges elements of the contained iterators using /// the ordering function. /// /// [`IntoIterator`] enabled version of [`Itertools::kmerge`]. /// /// ``` /// use itertools::kmerge; /// /// for elt in kmerge(vec![vec![0, 2, 4], vec![1, 3, 5], vec![6, 7]]) { /// /* loop body */ /// } /// ``` pub fn kmerge(iterable: I) -> KMerge<::IntoIter> where I: IntoIterator, I::Item: IntoIterator, <::Item as IntoIterator>::Item: PartialOrd, { kmerge_by(iterable, KMergeByLt) } /// An iterator adaptor that merges an abitrary number of base iterators /// according to an ordering function. /// /// Iterator element type is `I::Item`. /// /// See [`.kmerge_by()`](crate::Itertools::kmerge_by) for more /// information. #[must_use = "iterator adaptors are lazy and do nothing unless consumed"] pub struct KMergeBy where I: Iterator, { heap: Vec>, less_than: F, } impl fmt::Debug for KMergeBy where I: Iterator + fmt::Debug, I::Item: fmt::Debug, { debug_fmt_fields!(KMergeBy, heap); } /// Create an iterator that merges elements of the contained iterators. /// /// [`IntoIterator`] enabled version of [`Itertools::kmerge_by`]. pub fn kmerge_by( iterable: I, mut less_than: F, ) -> KMergeBy<::IntoIter, F> where I: IntoIterator, I::Item: IntoIterator, F: KMergePredicate<<::Item as IntoIterator>::Item>, { let iter = iterable.into_iter(); let (lower, _) = iter.size_hint(); let mut heap: Vec<_> = Vec::with_capacity(lower); heap.extend(iter.filter_map(|it| HeadTail::new(it.into_iter()))); heapify(&mut heap, |a, b| less_than.kmerge_pred(&a.head, &b.head)); KMergeBy { heap, less_than } } impl Clone for KMergeBy where I: Iterator + Clone, I::Item: Clone, F: Clone, { clone_fields!(heap, less_than); } impl Iterator for KMergeBy where I: Iterator, F: KMergePredicate, { type Item = I::Item; fn next(&mut self) -> Option { if self.heap.is_empty() { return None; } let result = if let Some(next) = self.heap[0].next() { next } else { self.heap.swap_remove(0).head }; let less_than = &mut self.less_than; sift_down(&mut self.heap, 0, |a, b| { less_than.kmerge_pred(&a.head, &b.head) }); Some(result) } fn size_hint(&self) -> (usize, Option) { #[allow(deprecated)] //TODO: once msrv hits 1.51. replace `fold1` with `reduce` self.heap .iter() .map(|i| i.size_hint()) .fold1(size_hint::add) .unwrap_or((0, Some(0))) } } impl FusedIterator for KMergeBy where I: Iterator, F: KMergePredicate, { }