#Skip to menu

Rust: Binary Search Tree

Published:

In this article we’ll see how to implement a Binary Search Tree in Rust. You can also get the final source code.

Node

A tree is composed of nodes. Each node has a key, a value and two sub-trees: a left one and a right one. We want to be able to store integers, strings, structs or any other data structure in our BST, so we use generics both for the key and the value. A node doesn’t always have sub-trees, and we use Option to represent the possibility of presence or absence. Here is the Node struct:

struct Node<K,V> {
    left: Option<Box<Node<K,V>>>,
    right: Option<Box<Node<K,V>>>,
    key: K,
    value: V,
}

Rust has many types of pointers, the most important ones are: references (&T) and Boxes.

  • References (&T) borrow a resource for a limited time. They don’t own the resource and are valid as long as the resource they point to is valid.
  • Boxes instead own the resource they point to. When the Box goes out of scope the resource is deallocated.

Here’s an example:

struct Node {
    field: u32,
}

fn main() {
    let x;
    {
        let y = Node{field: 3};
        x = &y;
    }
    // error: here x holds a reference to y
    //        but y goes out of scope
}

As you can see we can only reference a Node while we are inside its scope. What we need instead is to own the Node we refer to…

struct Node {
    field: u32,
}

fn main() {
    let x;
    {
        let y = Node{field: 3};
        x = Box::new(y);
    }
    // ok: x now owns y
}

Methods

We need a method to create a new node given a key and a value:

impl<K, V> Node<K, V> {
    fn new(k: K, v: V) -> Node<K, V> {
        Node {left: None,
              right: None,
              key: k,
              value: v,
        }
    }
}

BSTs are sorted which means that nodes on the left are less than the current node and those on the right are greater. How we decide if your generic type K is lesser or greater? We use a trait constrain which says that our key can be of any type as long as said type implements the Ord trait.

impl<K: Ord, V> Node<K, V> {

The insert method is defined recursively. If the Node to be inserted is less follow the left sub-tree, if it’s greater follow the right one, otherwise do nothing.

    fn insert(&mut self, n: Node<K, V>) {
        match n.key.cmp(&self.key) {
            Less => {
                match self.left {
                    None => self.left = Some(Box::new(n)),
                    Some(ref mut l) => l.insert(n),
                }
            },
            Greater => {
                match self.right {
                    None => self.right = Some(Box::new(n)),
                    Some(ref mut r) => r.insert(n),
                }
            },
            _ => {},
        }
    }

We do a similar thing with get.

    fn get(&self, key: &K) -> Option<&V> {
        match key.cmp(&self.key) {
            Equal => {
                Some(&self.value)
            },
            Less => {
                if let Some(ref l) = self.left {
                    l.get(key)
                } else {
                    None
                }
            },
            Greater => {
                if let Some(ref r) = self.right {
                    r.get(key)
                } else {
                    None
                }
            }
        }
    }

Instead of letting the user manipulate Nodes, we use another struct called Bst which holds the root of our tree and optionally other fields (like the length of the tree).

struct Bst<K,V> {
    root: Option<Box<Node<K,V>>>,
}
impl<K, V> Bst<K, V> {
    fn new() -> Bst<K, V> {
        Bst {root: None,
        }
    }
}

impl<K: Ord, V> Bst<K, V> {
    fn insert(&mut self, k: K, v: V) {
        match self.root {
            None => self.root = Some(Box::new(Node::new(k, v))),
            Some(ref mut r) => r.insert(Node::new(k, v)),
        }
    }

    fn get(&self, k: &K) -> Option<&V> {
        match self.root {
            None => None,
            Some(ref r) => r.get(k),
        }
    }
}

Here is an example

fn main() {
   let mut bst = Bst::new();
   bst.insert(2, "2");
   bst.insert(1, "1");
   bst.insert(3, "3");
   println!("{}", bst.get(2));
}

Traits

To make our BST more ergonomic to use, we need to implement some traits from the standard library.

Index

With this trait we can do bst[2] instead of bst.get(2).

impl<K: Ord, V> Index<K> for Bst<K, V>{
    type Output = V;

    fn index(&self, index: K) -> &V {
        self.get(&index).expect("no entry found for key")
    }
}

Iter

TODO

derive

Rust (like Haskell) can automaticaly implement some traits, when asked to do so, using the derive attribute. Add…

#[derive(Debug)]
struct Node<K,V> {

…and…

#[derive(Debug, Default)]
struct Bst<K, V> {

this will implement the Debug and Default traits.

Grapvhiz

Graphviz is a graph visualization software. We can generate a DOT file and pass the output to the dot program which will generate the graph of our tree.

impl<K: std::fmt::Display, V> Node<K, V> {
    fn dot_leaf(&self, leaf: &Option<Box<Node<K, V>>>, c: &mut u32) {
        match *leaf {
            None => {
                println!("null{} [shape=point];", c);
                println!("{} -> null{};", &self.key, c);
                *c += 1;
            },
            Some(ref l) => {
                println!("{} -> {}", &self.key, &l.key);
                l.dot(c);
            }
        }
    }

    fn dot(&self, c: &mut u32) {
        self.dot_leaf(&self.left, c);
        self.dot_leaf(&self.right, c);
    }
}

impl<K: std::fmt::Display, V> Bst<K, V> {
    fn dot(&self) {
        println!("digraph BST {{");

        let mut c = 0u32;

        match self.root {
            None => {},
            Some(ref r) => { r.dot(&mut c); },
        }

        println!("}}");
    }
}

Download the example and run it:

$ rustc dot.rs -o /tmp/dot && /tmp/dot | dot -Tx11

Source code

bst.rs dot.rs

Exercises for the reader

  • Implement a remove method.
  • Implement a length method on Bst. One way is to recursively calculate each time the length of the tree, the other is to have a length field in struct Bst and make insert return a true if a new node was inserted, false otherwise (the same for remove).
  • Implement a self-balancing tree (for example a Splay Tree).
  • Implement the IndexMut trait.

See what people say about this article.


Tags: , ,