7 /// Head element and Tail iterator pair
9 /// `PartialEq`, `Eq`, `PartialOrd` and `Ord` are implemented by comparing sequences based on
10 /// first items (which are guaranteed to exist).
12 /// The meanings of `PartialOrd` and `Ord` are reversed so as to turn the heap used in
13 /// `KMerge` into a min-heap.
25 /// Constructs a `HeadTail` from an `Iterator`. Returns `None` if the `Iterator` is empty.
26 fn new(mut it
: I
) -> Option
<HeadTail
<I
>> {
36 /// Get the next element and update `head`, returning the old head in `Some`.
38 /// Returns `None` when the tail is exhausted (only `head` then remains).
39 fn next(&mut self) -> Option
<I
::Item
> {
40 if let Some(next
) = self.tail
.next() {
41 Some(replace(&mut self.head
, next
))
47 /// Hints at the size of the sequence, same as the `Iterator` method.
48 fn size_hint(&self) -> (usize, Option
<usize>) {
49 size_hint
::add_scalar(self.tail
.size_hint(), 1)
53 impl<I
> Clone
for HeadTail
<I
>
54 where I
: Iterator
+ Clone
,
57 clone_fields
!(head
, tail
);
60 /// Make `data` a heap (min-heap w.r.t the sorting).
61 fn heapify
<T
, S
>(data
: &mut [T
], mut less_than
: S
)
62 where S
: FnMut(&T
, &T
) -> bool
64 for i
in (0..data
.len() / 2).rev() {
65 sift_down(data
, i
, &mut less_than
);
69 /// Sift down element at `index` (`heap` is a min-heap wrt the ordering)
70 fn sift_down
<T
, S
>(heap
: &mut [T
], index
: usize, mut less_than
: S
)
71 where S
: FnMut(&T
, &T
) -> bool
73 debug_assert
!(index
<= heap
.len());
75 let mut child
= 2 * pos
+ 1;
76 // the `pos` conditional is to avoid a bounds check
77 while pos
< heap
.len() && child
< heap
.len() {
78 let right
= child
+ 1;
80 // pick the smaller of the two children
81 if right
< heap
.len() && less_than(&heap
[right
], &heap
[child
]) {
85 // sift down is done if we are already in order
86 if !less_than(&heap
[child
], &heap
[pos
]) {
89 heap
.swap(pos
, child
);
95 /// An iterator adaptor that merges an abitrary number of base iterators in ascending order.
96 /// If all base iterators are sorted (ascending), the result is sorted.
98 /// Iterator element type is `I::Item`.
100 /// See [`.kmerge()`](../trait.Itertools.html#method.kmerge) for more information.
101 #[must_use = "iterator adaptors are lazy and do nothing unless consumed"]
102 pub type KMerge
<I
> = KMergeBy
<I
, KMergeByLt
>;
104 pub trait KMergePredicate
<T
> {
105 fn kmerge_pred(&mut self, a
: &T
, b
: &T
) -> bool
;
109 pub struct KMergeByLt
;
111 impl<T
: PartialOrd
> KMergePredicate
<T
> for KMergeByLt
{
112 fn kmerge_pred(&mut self, a
: &T
, b
: &T
) -> bool
{
117 impl<T
, F
: FnMut(&T
, &T
)->bool
> KMergePredicate
<T
> for F
{
118 fn kmerge_pred(&mut self, a
: &T
, b
: &T
) -> bool
{
123 /// Create an iterator that merges elements of the contained iterators using
124 /// the ordering function.
126 /// Equivalent to `iterable.into_iter().kmerge()`.
129 /// use itertools::kmerge;
131 /// for elt in kmerge(vec![vec![0, 2, 4], vec![1, 3, 5], vec![6, 7]]) {
135 pub fn kmerge
<I
>(iterable
: I
) -> KMerge
<<I
::Item
as IntoIterator
>::IntoIter
>
136 where I
: IntoIterator
,
137 I
::Item
: IntoIterator
,
138 <<I
as IntoIterator
>::Item
as IntoIterator
>::Item
: PartialOrd
140 kmerge_by(iterable
, KMergeByLt
)
143 /// An iterator adaptor that merges an abitrary number of base iterators
144 /// according to an ordering function.
146 /// Iterator element type is `I::Item`.
148 /// See [`.kmerge_by()`](../trait.Itertools.html#method.kmerge_by) for more
150 #[must_use = "iterator adaptors are lazy and do nothing unless consumed"]
151 pub struct KMergeBy
<I
, F
>
154 heap
: Vec
<HeadTail
<I
>>,
158 impl<I
, F
> fmt
::Debug
for KMergeBy
<I
, F
>
159 where I
: Iterator
+ fmt
::Debug
,
162 debug_fmt_fields
!(KMergeBy
, heap
);
165 /// Create an iterator that merges elements of the contained iterators.
167 /// Equivalent to `iterable.into_iter().kmerge_by(less_than)`.
168 pub fn kmerge_by
<I
, F
>(iterable
: I
, mut less_than
: F
)
169 -> KMergeBy
<<I
::Item
as IntoIterator
>::IntoIter
, F
>
170 where I
: IntoIterator
,
171 I
::Item
: IntoIterator
,
172 F
: KMergePredicate
<<<I
as IntoIterator
>::Item
as IntoIterator
>::Item
>,
174 let iter
= iterable
.into_iter();
175 let (lower
, _
) = iter
.size_hint();
176 let mut heap
: Vec
<_
> = Vec
::with_capacity(lower
);
177 heap
.extend(iter
.filter_map(|it
| HeadTail
::new(it
.into_iter())));
178 heapify(&mut heap
, |a
, b
| less_than
.kmerge_pred(&a
.head
, &b
.head
));
179 KMergeBy { heap, less_than }
182 impl<I
, F
> Clone
for KMergeBy
<I
, F
>
183 where I
: Iterator
+ Clone
,
187 clone_fields
!(heap
, less_than
);
190 impl<I
, F
> Iterator
for KMergeBy
<I
, F
>
192 F
: KMergePredicate
<I
::Item
>
196 fn next(&mut self) -> Option
<Self::Item
> {
197 if self.heap
.is_empty() {
200 let result
= if let Some(next
) = self.heap
[0].next() {
203 self.heap
.swap_remove(0).head
205 let less_than
= &mut self.less_than
;
206 sift_down(&mut self.heap
, 0, |a
, b
| less_than
.kmerge_pred(&a
.head
, &b
.head
));
210 fn size_hint(&self) -> (usize, Option
<usize>) {
212 .map(|i
| i
.size_hint())
213 .fold1(size_hint
::add
)
214 .unwrap_or((0, Some(0)))