melekes · December 27, 2021 09:44 · Dec 27, 2021
diff --git a/bst.rs b/bst.rs
@@ -0,0 +1,209 @@
+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());
+}
No results found