]>
Commit | Line | Data |
---|---|---|
f9f354fc XL |
1 | use crate::size_hint; |
2 | use crate::Itertools; | |
3 | ||
4 | use std::mem::replace; | |
5 | use std::fmt; | |
6 | ||
7 | /// Head element and Tail iterator pair | |
8 | /// | |
9 | /// `PartialEq`, `Eq`, `PartialOrd` and `Ord` are implemented by comparing sequences based on | |
10 | /// first items (which are guaranteed to exist). | |
11 | /// | |
12 | /// The meanings of `PartialOrd` and `Ord` are reversed so as to turn the heap used in | |
13 | /// `KMerge` into a min-heap. | |
14 | #[derive(Debug)] | |
15 | struct HeadTail<I> | |
16 | where I: Iterator | |
17 | { | |
18 | head: I::Item, | |
19 | tail: I, | |
20 | } | |
21 | ||
22 | impl<I> HeadTail<I> | |
23 | where I: Iterator | |
24 | { | |
25 | /// Constructs a `HeadTail` from an `Iterator`. Returns `None` if the `Iterator` is empty. | |
26 | fn new(mut it: I) -> Option<HeadTail<I>> { | |
27 | let head = it.next(); | |
28 | head.map(|h| { | |
29 | HeadTail { | |
30 | head: h, | |
31 | tail: it, | |
32 | } | |
33 | }) | |
34 | } | |
35 | ||
36 | /// Get the next element and update `head`, returning the old head in `Some`. | |
37 | /// | |
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)) | |
42 | } else { | |
43 | None | |
44 | } | |
45 | } | |
46 | ||
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) | |
50 | } | |
51 | } | |
52 | ||
53 | impl<I> Clone for HeadTail<I> | |
54 | where I: Iterator + Clone, | |
55 | I::Item: Clone | |
56 | { | |
57 | clone_fields!(head, tail); | |
58 | } | |
59 | ||
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 | |
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) | |
70 | fn sift_down<T, S>(heap: &mut [T], index: usize, mut less_than: S) | |
71 | where S: FnMut(&T, &T) -> bool | |
72 | { | |
73 | debug_assert!(index <= heap.len()); | |
74 | let mut pos = index; | |
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; | |
79 | ||
80 | // pick the smaller of the two children | |
81 | if right < heap.len() && less_than(&heap[right], &heap[child]) { | |
82 | child = right; | |
83 | } | |
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 | } | |
94 | ||
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. | |
97 | /// | |
98 | /// Iterator element type is `I::Item`. | |
99 | /// | |
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>; | |
103 | ||
104 | pub trait KMergePredicate<T> { | |
105 | fn kmerge_pred(&mut self, a: &T, b: &T) -> bool; | |
106 | } | |
107 | ||
108 | #[derive(Clone)] | |
109 | pub struct KMergeByLt; | |
110 | ||
111 | impl<T: PartialOrd> KMergePredicate<T> for KMergeByLt { | |
112 | fn kmerge_pred(&mut self, a: &T, b: &T) -> bool { | |
113 | a < b | |
114 | } | |
115 | } | |
116 | ||
117 | impl<T, F: FnMut(&T, &T)->bool> KMergePredicate<T> for F { | |
118 | fn kmerge_pred(&mut self, a: &T, b: &T) -> bool { | |
119 | self(a, b) | |
120 | } | |
121 | } | |
122 | ||
123 | /// Create an iterator that merges elements of the contained iterators using | |
124 | /// the ordering function. | |
125 | /// | |
126 | /// Equivalent to `iterable.into_iter().kmerge()`. | |
127 | /// | |
128 | /// ``` | |
129 | /// use itertools::kmerge; | |
130 | /// | |
131 | /// for elt in kmerge(vec![vec![0, 2, 4], vec![1, 3, 5], vec![6, 7]]) { | |
132 | /// /* loop body */ | |
133 | /// } | |
134 | /// ``` | |
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 | |
139 | { | |
140 | kmerge_by(iterable, KMergeByLt) | |
141 | } | |
142 | ||
143 | /// An iterator adaptor that merges an abitrary number of base iterators | |
144 | /// according to an ordering function. | |
145 | /// | |
146 | /// Iterator element type is `I::Item`. | |
147 | /// | |
148 | /// See [`.kmerge_by()`](../trait.Itertools.html#method.kmerge_by) for more | |
149 | /// information. | |
150 | #[must_use = "iterator adaptors are lazy and do nothing unless consumed"] | |
151 | pub struct KMergeBy<I, F> | |
152 | where I: Iterator, | |
153 | { | |
154 | heap: Vec<HeadTail<I>>, | |
155 | less_than: F, | |
156 | } | |
157 | ||
158 | impl<I, F> fmt::Debug for KMergeBy<I, F> | |
159 | where I: Iterator + fmt::Debug, | |
160 | I::Item: fmt::Debug, | |
161 | { | |
162 | debug_fmt_fields!(KMergeBy, heap); | |
163 | } | |
164 | ||
165 | /// Create an iterator that merges elements of the contained iterators. | |
166 | /// | |
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>, | |
173 | { | |
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 } | |
180 | } | |
181 | ||
182 | impl<I, F> Clone for KMergeBy<I, F> | |
183 | where I: Iterator + Clone, | |
184 | I::Item: Clone, | |
185 | F: Clone, | |
186 | { | |
187 | clone_fields!(heap, less_than); | |
188 | } | |
189 | ||
190 | impl<I, F> Iterator for KMergeBy<I, F> | |
191 | where I: Iterator, | |
192 | F: KMergePredicate<I::Item> | |
193 | { | |
194 | type Item = I::Item; | |
195 | ||
196 | fn next(&mut self) -> Option<Self::Item> { | |
197 | if self.heap.is_empty() { | |
198 | return None; | |
199 | } | |
200 | let result = if let Some(next) = self.heap[0].next() { | |
201 | next | |
202 | } else { | |
203 | self.heap.swap_remove(0).head | |
204 | }; | |
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)); | |
207 | Some(result) | |
208 | } | |
209 | ||
210 | fn size_hint(&self) -> (usize, Option<usize>) { | |
211 | self.heap.iter() | |
212 | .map(|i| i.size_hint()) | |
213 | .fold1(size_hint::add) | |
214 | .unwrap_or((0, Some(0))) | |
215 | } | |
216 | } |