mini-avl/src/set.rs

use crate::node::Node;

#[derive(Debug)]
pub struct Set<T: Ord> {
    root: Option<Box<Node<T>>>,
}

impl<T: Ord> Default for Set<T> {
    fn default() -> Self {
        Set { root: None }
    }
}

impl<T: Ord> Set<T> {
    pub fn new() -> Self {
        Default::default()
    }

    pub fn insert(&mut self, value: T) -> bool {
        if let Some(ref mut node) = self.root {
            node.insert(value)
        } else {
            self.root = Some(Box::new(Node::new(value)));
            true
        }
    }

    pub fn count(&self) -> i32 {
        self.root.as_ref().map_or(0, |n| n.count())
    }

    pub fn count_lt(&self, value: &T) -> i32 {
        self.count() - self.count_ge(value)
    }

    pub fn count_le(&self, value: &T) -> i32 {
        self.root.as_ref().map_or(0, |n| n.count_le(value))
    }

    pub fn count_ge(&self, value: &T) -> i32 {
        self.root.as_ref().map_or(0, |n| n.count_ge(value))
    }

    pub fn count_gt(&self, value: &T) -> i32 {
        self.count() - self.count_le(value)
    }
}

mod tests {
    #[test]
    fn test_count() {
        let mut tree = super::Set::new();
        for i in 1..11 {
            assert!(tree.insert(i));
            assert_eq!(tree.count(), i);
        }
        for i in 1..11 {
            assert_eq!(tree.count_lt(&i), i - 1);
            assert_eq!(tree.count_le(&i), i);
            assert_eq!(tree.count_ge(&i), 11 - i);
            assert_eq!(tree.count_gt(&i), 10 - i);
        }
    }
}