~cdv/tournament_tree

b01452ee767338aa0dd23b111fa9f93f53874281 — Chris Vittal 4 months ago 4efb736 main
Actually write the tournament tree
2 files changed, 314 insertions(+), 3 deletions(-)

M Cargo.toml
M src/lib.rs
M Cargo.toml => Cargo.toml +3 -0
@@ 1,8 1,11 @@
[package]
name = "tournament_tree"
description = "A tournament tree library"
version = "0.1.0"
authors = ["Chris Vittal <chris@vittal.dev>"]
edition = "2018"
repository = "https://git.sr.ht/~cdv/tournament_tree"
license = "MIT OR Apache-2.0"

# See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html


M src/lib.rs => src/lib.rs +311 -3
@@ 1,7 1,315 @@
//! A library implementing a tournament tree data structure.
//!
//! Tournament trees are, conceptually, complete binary trees holding the result
//! of comparisons between each element of the tree. They can be used as a fixed
//! size priority queue for applications like out-of-core sorting, or
//! implementing many way joins.
//!
//! ```
//! use tournament_tree::TournamentTree;
//! # use std::cmp::Reverse;
//! let mut data1 = vec![3, 1, 4, 1, 5, 2, 6, 5];
//! let tourney_tree = TournamentTree::from(data1.clone());
//! data1.sort_by_key(|&e| Reverse(e));
//! let data2: Vec<_> = tourney_tree.into_iter().collect();
//! assert_eq!(data1, data2);
//! ```
use std::{iter::FromIterator, ops::Deref};

/// A tournament tree. See the module documentation for more details.
///
/// This is a max tree.
#[derive(Debug, Clone)]
pub struct TournamentTree<T> {
    data: Vec<T>,
    tree: Vec<usize>,
}

/// Immutable tournament tree iterator.
#[derive(Debug, Clone)]
pub struct Iter<'tree, T> {
    data: &'tree TournamentTree<T>,
    tree: Vec<usize>,
}

/// An iterator that moves out of the tournament
/// tree.
#[derive(Debug, Clone)]
pub struct IntoIter<T> {
    data: Vec<Option<T>>,
    tree: Vec<usize>,
}

impl<T: Ord> TournamentTree<T> {
    /// Returns the winner.
    ///
    /// # Panics
    /// If the tree is empty.
    pub fn winner(&self) -> &T {
        &self[self.winner_idx()]
    }

    /// The Index of the winner.
    ///
    /// # Panics
    /// If the tree is empty.
    pub fn winner_idx(&self) -> usize {
        self.tree[0]
    }

    fn build(data: Vec<T>) -> Self {
        let mut tree = vec![usize::MAX; data.len()];
        let k = data.len();
        for i in 0..data.len() {
            let mut winner = i;
            let mut j = (i + k) / 2;
            while j != 0 && tree[j] != usize::MAX {
                let challenger = tree[j];
                if data[challenger] > data[winner] {
                    tree[j] = winner;
                    winner = challenger;
                }
                j = j / 2;
            }
            tree[j] = winner;
        }

        Self { data, tree }
    }

    pub fn iter(&self) -> Iter<T> {
        Iter {
            data: self,
            tree: self.tree.clone(),
        }
    }
}

impl<T: Ord> Deref for TournamentTree<T> {
    type Target = [T];

    fn deref(&self) -> &Self::Target {
        &self.data
    }
}

impl<T: Ord> From<Vec<T>> for TournamentTree<T> {
    fn from(data: Vec<T>) -> Self {
        Self::build(data)
    }
}

impl<T: Ord> Into<Vec<T>> for TournamentTree<T> {
    fn into(self) -> Vec<T> {
        self.data
    }
}

impl<T: Ord> AsRef<[T]> for TournamentTree<T> {
    fn as_ref(&self) -> &[T] {
        &self.data
    }
}

impl<T: Ord> FromIterator<T> for TournamentTree<T> {
    fn from_iter<I: IntoIterator<Item = T>>(iter: I) -> Self {
        let data = Vec::from_iter(iter);
        Self::build(data)
    }
}

impl<T: Ord> IntoIterator for TournamentTree<T> {
    type Item = T;
    type IntoIter = IntoIter<T>;
    fn into_iter(self) -> Self::IntoIter {
        IntoIter {
            data: self.data.into_iter().map(Some).collect(),
            tree: self.tree,
        }
    }
}

impl<'tree, T: Ord> IntoIterator for &'tree TournamentTree<T> {
    type Item = &'tree T;
    type IntoIter = Iter<'tree, T>;

    fn into_iter(self) -> Self::IntoIter {
        self.iter()
    }
}

impl<T: Ord> Iterator for IntoIter<T> {
    type Item = T;

    fn next(&mut self) -> Option<Self::Item> {
        let k = self.data.len();
        if self.tree[0] == k {
            return None;
        }

        let i = self.tree[0];
        let mut winner = k;
        let mut j = (i + k) / 2;
        while j != 0 {
            let challenger = self.tree[j];
            if challenger != k && (winner == k || self.data[challenger] > self.data[winner]) {
                self.tree[j] = winner;
                winner = challenger;
            }
            j = j / 2;
        }
        self.tree[j] = winner;

        self.data[i].take()
    }
}

impl<'tree, T: Ord> Iterator for Iter<'tree, T> {
    type Item = &'tree T;

    fn next(&mut self) -> Option<Self::Item> {
        let k = self.data.len();
        if self.tree[0] == k {
            return None;
        }

        let i = self.tree[0];
        let mut winner = k;
        let mut j = (i + k) / 2;
        while j != 0 {
            let challenger = self.tree[j];
            if challenger != k
                && (winner == k || self.data.data[challenger] > self.data.data[winner])
            {
                self.tree[j] = winner;
                winner = challenger;
            }
            j = j / 2;
        }
        self.tree[j] = winner;

        Some(&self.data.data[i])
    }
}

/// A bounded priority queue using a `TournamentTree` for a backing structure.
///
/// Will replace each element removed with the next one from it's iterator until
/// all iterators are exhausted. If the iterators iterate in reverse sorted order, the
/// output will be in sorted order.
pub struct TournamentTreeIterator<T> {
    tt: TournamentTree<Option<T>>,
    its: Vec<Box<dyn Iterator<Item = T>>>,
}

pub struct TournamentTreeIteratorBuilder<T> {
    its: Vec<Box<dyn Iterator<Item = T>>>,
}

impl<T: Ord> TournamentTreeIterator<T> {
    /// Creates a new `TournamentTreeIteratorBuilder` for constructing a
    /// `TournamentTreeIterator`
    pub fn builder() -> TournamentTreeIteratorBuilder<T> {
        TournamentTreeIteratorBuilder { its: Vec::new() }
    }
}

impl<T: Ord> TournamentTreeIterator<T> {
    fn adv(&mut self) -> Option<T> {
        let i = self.tt.winner_idx();
        let k = self.tt.data.len();
        let ret = self.tt.data[i].take();
        self.tt.data[i] = self.its[i].next();
        let mut winner = i;
        let mut j = (i + k) / 2;
        while j != 0 {
            let challenger = self.tt.tree[j];
            if self.tt[challenger] > self.tt[winner] {
                self.tt.tree[j] = winner;
                winner = challenger;
            }
            j = j / 2;
        }
        self.tt.tree[j] = winner;

        ret
    }
}

impl<T: Ord> Iterator for TournamentTreeIterator<T> {
    type Item = T;

    fn next(&mut self) -> Option<Self::Item> {
        if self.tt.winner().is_none() {
            return None;
        }
        self.adv()
    }
}

impl<T: Ord> TournamentTreeIteratorBuilder<T> {
    /// Add an iterator to the eventual `TournamentTreeIterator`
    pub fn add(&mut self, it: impl Iterator<Item = T> + 'static) -> &mut Self {
        self.its.push(Box::new(it));
        self
    }

    /// Consumes the builder, returning the `TournamentTreeIterator`.
    pub fn build(mut self) -> TournamentTreeIterator<T> {
        let mut data = Vec::with_capacity(self.its.len());
        for it in self.its.iter_mut() {
            data.push(it.next());
        }

        TournamentTreeIterator {
            tt: TournamentTree::from(data),
            its: self.its,
        }
    }
}

#[cfg(test)]
mod tests {
mod test {
    use super::*;

    #[test]
    fn basic() {
        let data = vec![3, 1, 4, 1, 5, 2, 6, 5];
        let tourney_tree = TournamentTree::build(data);
        assert_eq!(*tourney_tree.winner(), 6);
    }

    #[test]
    fn iter() {
        use std::cmp::Reverse;
        let mut data1 = vec![3, 1, 4, 1, 5, 2, 6, 5];
        let tourney_tree = TournamentTree::build(data1.clone());
        data1.sort_by_key(|&e| Reverse(e));
        let data2: Vec<_> = tourney_tree.iter().copied().collect();
        assert_eq!(data1, data2);
    }

    #[test]
    fn into_iter() {
        use std::cmp::Reverse;
        let mut data1 = vec![3, 1, 4, 1, 5, 2, 6, 5];
        let tourney_tree = TournamentTree::build(data1.clone());
        data1.sort_by_key(|&e| Reverse(e));
        let data2: Vec<_> = tourney_tree.into_iter().collect();
        assert_eq!(data1, data2);
    }

    #[test]
    fn it_works() {
        assert_eq!(2 + 2, 4);
    fn iterator() {
        let mut ttib = TournamentTreeIterator::builder();
        ttib.add(vec![7, 5, 3, 1].into_iter())
            .add(vec![16, 8, 4, 2, 1].into_iter())
            .add(vec![11, 7, 5, 3, 2].into_iter())
            .add(vec![25, 16, 9, 4, 1, 0].into_iter());
        let v: Vec<_> = ttib.build().collect();
        assert_eq!(
            v,
            [25, 16, 16, 11, 9, 8, 7, 7, 5, 5, 4, 4, 3, 3, 2, 2, 1, 1, 1, 0]
        );
    }
}