tournament_tree/
lib.rs

1//! A library implementing a tournament tree data structure.
2//!
3//! Tournament trees are, conceptually, complete binary trees holding the result
4//! of comparisons between each element of the tree. They can be used as a fixed
5//! size priority queue for applications like out-of-core sorting, or
6//! implementing many way joins.
7//!
8//! ```
9//! use tournament_tree::TournamentTree;
10//! # use std::cmp::Reverse;
11//! let mut data1 = vec![3, 1, 4, 1, 5, 2, 6, 5];
12//! let tourney_tree = TournamentTree::from(data1.clone());
13//! data1.sort_by_key(|&e| Reverse(e));
14//! let data2: Vec<_> = tourney_tree.into_iter().collect();
15//! assert_eq!(data1, data2);
16//! ```
17use std::{iter::FromIterator, ops::Deref};
18
19/// A tournament tree. See the module documentation for more details.
20///
21/// This is a max tree.
22#[derive(Debug, Clone)]
23pub struct TournamentTree<T> {
24    data: Vec<T>,
25    tree: Vec<usize>,
26}
27
28/// Immutable tournament tree iterator.
29#[derive(Debug, Clone)]
30pub struct Iter<'tree, T> {
31    data: &'tree TournamentTree<T>,
32    tree: Vec<usize>,
33}
34
35/// An iterator that moves out of the tournament
36/// tree.
37#[derive(Debug, Clone)]
38pub struct IntoIter<T> {
39    data: Vec<Option<T>>,
40    tree: Vec<usize>,
41}
42
43impl<T: Ord> TournamentTree<T> {
44    /// Returns the winner.
45    ///
46    /// # Panics
47    /// If the tree is empty.
48    pub fn winner(&self) -> &T {
49        &self[self.winner_idx()]
50    }
51
52    /// The Index of the winner.
53    ///
54    /// # Panics
55    /// If the tree is empty.
56    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
195/// A bounded priority queue using a `TournamentTree` for a backing structure.
196///
197/// Will replace each element removed with the next one from it's iterator until
198/// all iterators are exhausted. If the iterators iterate in reverse sorted order, the
199/// output will be in sorted order.
200pub 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    /// Creates a new `TournamentTreeIteratorBuilder` for constructing a
211    /// `TournamentTreeIterator`
212    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    /// Add an iterator to the eventual `TournamentTreeIterator`
252    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    /// Consumes the builder, returning the `TournamentTreeIterator`.
258    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}