1use std::{iter::FromIterator, ops::Deref};
18
19#[derive(Debug, Clone)]
23pub struct TournamentTree<T> {
24 data: Vec<T>,
25 tree: Vec<usize>,
26}
27
28#[derive(Debug, Clone)]
30pub struct Iter<'tree, T> {
31 data: &'tree TournamentTree<T>,
32 tree: Vec<usize>,
33}
34
35#[derive(Debug, Clone)]
38pub struct IntoIter<T> {
39 data: Vec<Option<T>>,
40 tree: Vec<usize>,
41}
42
43impl<T: Ord> TournamentTree<T> {
44 pub fn winner(&self) -> &T {
49 &self[self.winner_idx()]
50 }
51
52 pub fn winner_idx(&self) -> usize {
57 self.tree[0]
58 }
59
60 fn build(data: Vec<T>) -> Self {
61 let mut tree = vec![usize::MAX; data.len()];
62 let k = data.len();
63 for i in 0..data.len() {
64 let mut winner = i;
65 let mut j = (i + k) / 2;
66 while j != 0 && tree[j] != usize::MAX {
67 let challenger = tree[j];
68 if data[challenger] > data[winner] {
69 tree[j] = winner;
70 winner = challenger;
71 }
72 j = j / 2;
73 }
74 tree[j] = winner;
75 }
76
77 Self { data, tree }
78 }
79
80 pub fn iter(&self) -> Iter<T> {
81 Iter {
82 data: self,
83 tree: self.tree.clone(),
84 }
85 }
86}
87
88impl<T: Ord> Deref for TournamentTree<T> {
89 type Target = [T];
90
91 fn deref(&self) -> &Self::Target {
92 &self.data
93 }
94}
95
96impl<T: Ord> From<Vec<T>> for TournamentTree<T> {
97 fn from(data: Vec<T>) -> Self {
98 Self::build(data)
99 }
100}
101
102impl<T: Ord> Into<Vec<T>> for TournamentTree<T> {
103 fn into(self) -> Vec<T> {
104 self.data
105 }
106}
107
108impl<T: Ord> AsRef<[T]> for TournamentTree<T> {
109 fn as_ref(&self) -> &[T] {
110 &self.data
111 }
112}
113
114impl<T: Ord> FromIterator<T> for TournamentTree<T> {
115 fn from_iter<I: IntoIterator<Item = T>>(iter: I) -> Self {
116 let data = Vec::from_iter(iter);
117 Self::build(data)
118 }
119}
120
121impl<T: Ord> IntoIterator for TournamentTree<T> {
122 type Item = T;
123 type IntoIter = IntoIter<T>;
124 fn into_iter(self) -> Self::IntoIter {
125 IntoIter {
126 data: self.data.into_iter().map(Some).collect(),
127 tree: self.tree,
128 }
129 }
130}
131
132impl<'tree, T: Ord> IntoIterator for &'tree TournamentTree<T> {
133 type Item = &'tree T;
134 type IntoIter = Iter<'tree, T>;
135
136 fn into_iter(self) -> Self::IntoIter {
137 self.iter()
138 }
139}
140
141impl<T: Ord> Iterator for IntoIter<T> {
142 type Item = T;
143
144 fn next(&mut self) -> Option<Self::Item> {
145 let k = self.data.len();
146 if self.tree[0] == k {
147 return None;
148 }
149
150 let i = self.tree[0];
151 let mut winner = k;
152 let mut j = (i + k) / 2;
153 while j != 0 {
154 let challenger = self.tree[j];
155 if challenger != k && (winner == k || self.data[challenger] > self.data[winner]) {
156 self.tree[j] = winner;
157 winner = challenger;
158 }
159 j = j / 2;
160 }
161 self.tree[j] = winner;
162
163 self.data[i].take()
164 }
165}
166
167impl<'tree, T: Ord> Iterator for Iter<'tree, T> {
168 type Item = &'tree T;
169
170 fn next(&mut self) -> Option<Self::Item> {
171 let k = self.data.len();
172 if self.tree[0] == k {
173 return None;
174 }
175
176 let i = self.tree[0];
177 let mut winner = k;
178 let mut j = (i + k) / 2;
179 while j != 0 {
180 let challenger = self.tree[j];
181 if challenger != k
182 && (winner == k || self.data.data[challenger] > self.data.data[winner])
183 {
184 self.tree[j] = winner;
185 winner = challenger;
186 }
187 j = j / 2;
188 }
189 self.tree[j] = winner;
190
191 Some(&self.data.data[i])
192 }
193}
194
195pub struct TournamentTreeIterator<T> {
201 tt: TournamentTree<Option<T>>,
202 its: Vec<Box<dyn Iterator<Item = T>>>,
203}
204
205pub struct TournamentTreeIteratorBuilder<T> {
206 its: Vec<Box<dyn Iterator<Item = T>>>,
207}
208
209impl<T: Ord> TournamentTreeIterator<T> {
210 pub fn builder() -> TournamentTreeIteratorBuilder<T> {
213 TournamentTreeIteratorBuilder { its: Vec::new() }
214 }
215}
216
217impl<T: Ord> TournamentTreeIterator<T> {
218 fn adv(&mut self) -> Option<T> {
219 let i = self.tt.winner_idx();
220 let k = self.tt.data.len();
221 let ret = self.tt.data[i].take();
222 self.tt.data[i] = self.its[i].next();
223 let mut winner = i;
224 let mut j = (i + k) / 2;
225 while j != 0 {
226 let challenger = self.tt.tree[j];
227 if self.tt[challenger] > self.tt[winner] {
228 self.tt.tree[j] = winner;
229 winner = challenger;
230 }
231 j = j / 2;
232 }
233 self.tt.tree[j] = winner;
234
235 ret
236 }
237}
238
239impl<T: Ord> Iterator for TournamentTreeIterator<T> {
240 type Item = T;
241
242 fn next(&mut self) -> Option<Self::Item> {
243 if self.tt.winner().is_none() {
244 return None;
245 }
246 self.adv()
247 }
248}
249
250impl<T: Ord> TournamentTreeIteratorBuilder<T> {
251 pub fn add(&mut self, it: impl Iterator<Item = T> + 'static) -> &mut Self {
253 self.its.push(Box::new(it));
254 self
255 }
256
257 pub fn build(mut self) -> TournamentTreeIterator<T> {
259 let mut data = Vec::with_capacity(self.its.len());
260 for it in self.its.iter_mut() {
261 data.push(it.next());
262 }
263
264 TournamentTreeIterator {
265 tt: TournamentTree::from(data),
266 its: self.its,
267 }
268 }
269}
270
271#[cfg(test)]
272mod test {
273 use super::*;
274
275 #[test]
276 fn basic() {
277 let data = vec![3, 1, 4, 1, 5, 2, 6, 5];
278 let tourney_tree = TournamentTree::build(data);
279 assert_eq!(*tourney_tree.winner(), 6);
280 }
281
282 #[test]
283 fn iter() {
284 use std::cmp::Reverse;
285 let mut data1 = vec![3, 1, 4, 1, 5, 2, 6, 5];
286 let tourney_tree = TournamentTree::build(data1.clone());
287 data1.sort_by_key(|&e| Reverse(e));
288 let data2: Vec<_> = tourney_tree.iter().copied().collect();
289 assert_eq!(data1, data2);
290 }
291
292 #[test]
293 fn into_iter() {
294 use std::cmp::Reverse;
295 let mut data1 = vec![3, 1, 4, 1, 5, 2, 6, 5];
296 let tourney_tree = TournamentTree::build(data1.clone());
297 data1.sort_by_key(|&e| Reverse(e));
298 let data2: Vec<_> = tourney_tree.into_iter().collect();
299 assert_eq!(data1, data2);
300 }
301
302 #[test]
303 fn iterator() {
304 let mut ttib = TournamentTreeIterator::builder();
305 ttib.add(vec![7, 5, 3, 1].into_iter())
306 .add(vec![16, 8, 4, 2, 1].into_iter())
307 .add(vec![11, 7, 5, 3, 2].into_iter())
308 .add(vec![25, 16, 9, 4, 1, 0].into_iter());
309 let v: Vec<_> = ttib.build().collect();
310 assert_eq!(
311 v,
312 [25, 16, 16, 11, 9, 8, 7, 7, 5, 5, 4, 4, 3, 3, 2, 2, 1, 1, 1, 0]
313 );
314 }
315}