BST || k_largest element

bst.rs

use anyhow::anyhow;
use anyhow::Result;

use std::mem::swap;

type Child = Option<Box<Node>>;

#[derive(Debug, Eq, PartialOrd, PartialEq, Clone)]
pub struct Node {
    pub value: i32,
    pub left: Child,
    pub right: Child,
    pub size: usize,
}

impl Node {
    fn new(value: i32) -> Self {
        Node {
            value,
            left: None,
            right: None,
            size: 1,
        }
    }

    fn push(&mut self, value: i32) {
        if value > self.value {
            match self.right {
                None => swap(&mut self.right, &mut Some(Box::from(Node::new(value)))),
                Some(ref mut n) => n.push(value),
            }
        } else {
            match self.left {
                None => swap(&mut self.left, &mut Some(Box::from(Node::new(value)))),
                Some(ref mut n) => n.push(value),
            }
        }
        self.size += 1;
    }

    fn k_largest(&self, k: usize) -> Result<i32> {
        if k == 0 || k > self.size {
            return Err(anyhow!(
                "k must be within [1, {}], but given {}",
                self.size,
                k
            ));
        }

        // perform reverse search while keeping the count of visited nodes
        let (value, _) = Node::reverse_search(self, k, 1);

        value.ok_or(anyhow!("not found"))
    }

    fn largest(&self) -> (&Node, Vec<&Node>) {
        let mut largest = self;
        let mut ancestors = Vec::new();
        loop {
            match largest.right {
                None => break,
                Some(ref node) => {
                    ancestors.push(largest);
                    largest = node;
                }
            }
        }
        (largest, ancestors)
    }

    fn reverse_search(root: &Node, k: usize, num_visited: usize) -> (Option<i32>, usize) {
        // go to the bottom right element (can be root)
        let (largest, ancestors) = root.largest();

        if num_visited == k {
            return (Some(largest.value), num_visited);
        }

        let mut new_num_visited = num_visited;

        // search left child subtree (if any)
        if let Some(ref n) = largest.left {
            let (klargest, nn) = Node::reverse_search(n, k, num_visited + 1);
            new_num_visited = nn;
            if klargest.is_some() {
                return (klargest, new_num_visited);
            }
        }

        // go through all ancestors
        for parent in ancestors.iter().rev() {
            new_num_visited += 1;

            if new_num_visited == k {
                return (Some(parent.value), new_num_visited);
            }

            // go to parent's left subtree
            if let Some(ref left) = parent.left {
                let (klargest, new_num_visited) =
                    Node::reverse_search(left, k, new_num_visited + 1);
                if klargest.is_some() {
                    return (klargest, new_num_visited);
                }
            }
        }

        (None, new_num_visited)
    }
}

/// Binary search tree <https://en.wikipedia.org/wiki/Binary_search_tree>
#[derive(Debug, PartialOrd, PartialEq, Clone)]
pub struct Tree {
    root: Child,
}

impl Tree {
    /// Returns an empty tree.
    pub fn new() -> Self {
        Tree { root: None }
    }

    /// Adds a value into a tree.
    ///
    /// # Examples
    ///
    /// ```
    /// let mut tree = Tree::new();
    /// tree.push(1);
    /// tree.push(2);
    /// tree.push(3);
    /// ```
    pub fn push(&mut self, value: i32) {
        match self.root {
            None => {
                swap(&mut self.root, &mut Some(Box::from(Node::new(value))));
            }
            Some(ref mut n) => {
                n.push(value);
            }
        }
    }

    /// Returns the k (k -> [1; N] where N is the number of elements) largest value.
    ///
    /// # Examples
    ///
    /// ```
    /// let mut tree = Tree::new();
    /// tree.push(1);
    /// tree.push(2);
    /// tree.push(3);
    ///
    /// assert_eq!(3, tree.k_largest(1).unwrap());
    /// ```
    pub fn k_largest(&self, k: usize) -> Result<i32> {
        match self.root {
            None => Err(anyhow!("tree is empty")),
            Some(ref n) => n.k_largest(k),
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    #[test]
    fn k_largest_returns_correct_element() {
        let mut tree = Tree::new();
        tree.push(3);
        tree.push(2);
        tree.push(1);

        assert_eq!(3, tree.k_largest(1).unwrap());
        assert_eq!(2, tree.k_largest(2).unwrap());
        assert_eq!(1, tree.k_largest(3).unwrap());

        let mut tree = Tree::new();
        tree.push(3);
        tree.push(2);
        tree.push(1);
        tree.push(5);
        tree.push(8);
        tree.push(7);

        assert_eq!(8, tree.k_largest(1).unwrap());
        assert_eq!(7, tree.k_largest(2).unwrap());
        assert_eq!(5, tree.k_largest(3).unwrap());
        assert_eq!(3, tree.k_largest(4).unwrap());
        assert_eq!(2, tree.k_largest(5).unwrap());
        assert_eq!(1, tree.k_largest(6).unwrap());
    }
}

fn main() {
    let mut tree = Tree::new();
    tree.push(6);
    tree.push(2);
    tree.push(3);
    tree.push(4);
    tree.push(10);

    println!("{}", tree.k_largest(1).unwrap());
    println!("{}", tree.k_largest(2).unwrap());
    println!("{}", tree.k_largest(3).unwrap());
    println!("{}", tree.k_largest(4).unwrap());
    println!("{}", tree.k_largest(5).unwrap());
}