]>
Commit | Line | Data |
---|---|---|
c295e0f8 XL |
1 | use crate::size_hint; |
2 | use crate::Itertools; | |
3 | ||
4 | use alloc::vec::Vec; | |
5 | use std::iter::FusedIterator; | |
6 | use std::mem::replace; | |
7 | use 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)] | |
17 | struct HeadTail<I> | |
18 | where I: Iterator | |
19 | { | |
20 | head: I::Item, | |
21 | tail: I, | |
22 | } | |
23 | ||
24 | impl<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 | ||
55 | impl<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). | |
63 | fn 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) | |
72 | fn 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 | |
2b03887a | 83 | // use arithmetic to avoid an unpredictable branch |
c295e0f8 XL |
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. | |
c295e0f8 XL |
107 | pub type KMerge<I> = KMergeBy<I, KMergeByLt>; |
108 | ||
109 | pub trait KMergePredicate<T> { | |
110 | fn kmerge_pred(&mut self, a: &T, b: &T) -> bool; | |
111 | } | |
112 | ||
a2a8927a | 113 | #[derive(Clone, Debug)] |
c295e0f8 XL |
114 | pub struct KMergeByLt; |
115 | ||
116 | impl<T: PartialOrd> KMergePredicate<T> for KMergeByLt { | |
117 | fn kmerge_pred(&mut self, a: &T, b: &T) -> bool { | |
118 | a < b | |
119 | } | |
120 | } | |
121 | ||
122 | impl<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 | /// | |
2b03887a | 131 | /// [`IntoIterator`] enabled version of [`Itertools::kmerge`]. |
c295e0f8 XL |
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 | /// ``` | |
140 | pub 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"] | |
156 | pub struct KMergeBy<I, F> | |
157 | where I: Iterator, | |
158 | { | |
159 | heap: Vec<HeadTail<I>>, | |
160 | less_than: F, | |
161 | } | |
162 | ||
163 | impl<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 | /// | |
2b03887a | 172 | /// [`IntoIterator`] enabled version of [`Itertools::kmerge_by`]. |
c295e0f8 XL |
173 | pub 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 | ||
187 | impl<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 | ||
195 | impl<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>) { | |
2b03887a | 216 | #[allow(deprecated)] //TODO: once msrv hits 1.51. replace `fold1` with `reduce` |
c295e0f8 XL |
217 | self.heap.iter() |
218 | .map(|i| i.size_hint()) | |
219 | .fold1(size_hint::add) | |
220 | .unwrap_or((0, Some(0))) | |
221 | } | |
222 | } | |
223 | ||
224 | impl<I, F> FusedIterator for KMergeBy<I, F> | |
225 | where I: Iterator, | |
226 | F: KMergePredicate<I::Item> | |
227 | {} |