263 Advanced

AVL Tree — Self-Balancing BST

Trees

Tutorial Video

Text description (accessibility)

This video demonstrates the "AVL Tree — Self-Balancing BST" functional Rust example. Difficulty level: Advanced. Key concepts covered: Trees. Implement an AVL tree — a self-balancing binary search tree where the heights of left and right subtrees differ by at most 1. Key difference from OCaml: 1. **Move semantics:** Rotations in Rust consume the tree (`self` by value), making restructuring explicit; OCaml creates new nodes implicitly

Tutorial

The Problem

Implement an AVL tree — a self-balancing binary search tree where the heights of left and right subtrees differ by at most 1. The tree automatically rebalances via rotations after each insert.

🎯 Learning Outcomes

• Implementing complex recursive data structures with named enum fields

• Translating OCaml's nested pattern matching on constructors to Rust

• Using Box ownership transfer for tree rotations (consuming and restructuring)

• Maintaining invariants (balance factor) in persistent data structures

🦀 The Rust Way

Rust uses named struct fields in the Node variant for clarity. Rotations consume self (move semantics) and reconstruct the tree, which naturally expresses the restructuring. The Ord trait provides generic comparison.

Code Example

#[derive(Debug, Clone, PartialEq)]
pub enum Avl<T> {
    Empty,
    Node { left: Box<Avl<T>>, value: T, right: Box<Avl<T>>, height: i32 },
}

impl<T: Ord + Clone> Avl<T> {
    fn rotate_right(self) -> Self {
        match self {
            Avl::Node { left, value, right, .. } => match *left {
                Avl::Node { left: ll, value: lv, right: lr, .. } =>
                    Self::node(Self::node(*ll, lv, *lr), value, *right),
                _ => Self::node(*left, value, *right),
            },
            other => other,
        }
    }

    pub fn insert(&self, x: T) -> Self {
        match self {
            Avl::Empty => Self::node(Avl::Empty, x, Avl::Empty),
            Avl::Node { left, value, right, .. } => match x.cmp(value) {
                Ordering::Less => Self::node(left.insert(x), value.clone(), (**right).clone()).rebalance(),
                Ordering::Greater => Self::node((**left).clone(), value.clone(), right.insert(x)).rebalance(),
                Ordering::Equal => self.clone(),
            },
        }
    }
}

type 'a avl = Empty | Node of 'a avl * 'a * 'a avl * int

let height = function Empty -> 0 | Node (_, _, _, h) -> h
let node l v r = Node (l, v, r, 1 + max (height l) (height r))

let rotate_right = function
  | Node (Node (ll, lv, lr, _), v, r, _) -> node (node ll lv lr) v r
  | t -> t

let rec insert x = function
  | Empty -> node Empty x Empty
  | Node (l, v, r, _) ->
    if x < v then rebalance (node (insert x l) v r)
    else if x > v then rebalance (node l v (insert x r))
    else node l v r

Key Differences

Move semantics: Rotations in Rust consume the tree (self by value), making restructuring explicit; OCaml creates new nodes implicitly

Named fields: Rust's Node { left, value, right, height } is more readable than OCaml's positional tuple

Nested destructuring: OCaml matches two levels in one arm; Rust needs nested match blocks

Height caching: Both store height in nodes; Rust's i32 vs OCaml's int — same idea, explicit type

OCaml Approach

OCaml stores the height in each node: Node of 'a avl * 'a * 'a avl * int. Rotations are pattern matches that destructure two levels of the tree. The rebalance function checks the balance factor and applies the appropriate rotation.

Full Source

#![allow(clippy::all)]
//! AVL Tree — Self-Balancing BST
//!
//! OCaml: `type 'a avl = Empty | Node of 'a avl * 'a * 'a avl * int`
//! Rust: `enum Avl<T> { Empty, Node { left, value, right, height } }`
//!
//! An AVL tree maintains a balance invariant: for every node, the height
//! difference between left and right subtrees is at most 1. Rotations
//! restore balance after each insert.
use std::cmp::{max, Ordering};

/// Persistent AVL tree — each insert returns a new balanced tree.
#[derive(Debug, Clone, PartialEq)]
pub enum Avl<T> {
    Empty,
    Node {
        left: Box<Avl<T>>,
        value: T,
        right: Box<Avl<T>>,
        height: i32,
    },
}

impl<T: Ord + Clone> Avl<T> {
    /// Creates an empty AVL tree.
    pub fn empty() -> Self {
        Avl::Empty
    }

    /// Returns the height of the tree.
    pub fn height(&self) -> i32 {
        match self {
            Avl::Empty => 0,
            Avl::Node { height, .. } => *height,
        }
    }

    /// Creates a node, computing height from children.
    /// OCaml: `let node l v r = Node (l, v, r, 1 + max (height l) (height r))`
    fn node(left: Avl<T>, value: T, right: Avl<T>) -> Self {
        let h = 1 + max(left.height(), right.height());
        Avl::Node {
            left: Box::new(left),
            value,
            right: Box::new(right),
            height: h,
        }
    }

    /// Computes the balance factor (left height - right height).
    fn balance_factor(&self) -> i32 {
        match self {
            Avl::Empty => 0,
            Avl::Node { left, right, .. } => left.height() - right.height(),
        }
    }

    /// Right rotation for left-heavy trees.
    ///
    /// ```text
    ///     v              lv
    ///    / \            /  \
    ///   lv  r   =>    ll   v
    ///  / \                / \
    /// ll  lr             lr  r
    /// ```
    fn rotate_right(self) -> Self {
        match self {
            Avl::Node {
                left, value, right, ..
            } => match *left {
                Avl::Node {
                    left: ll,
                    value: lv,
                    right: lr,
                    ..
                } => Self::node(*ll, lv, Self::node(*lr, value, *right)),
                _ => Self::node(*left, value, *right),
            },
            other => other,
        }
    }

    /// Left rotation for right-heavy trees.
    ///
    /// ```text
    ///   v                rv
    ///  / \              /  \
    /// l   rv    =>     v    rr
    ///    / \          / \
    ///   rl  rr      l   rl
    /// ```
    fn rotate_left(self) -> Self {
        match self {
            Avl::Node {
                left, value, right, ..
            } => match *right {
                Avl::Node {
                    left: rl,
                    value: rv,
                    right: rr,
                    ..
                } => Self::node(Self::node(*left, value, *rl), rv, *rr),
                _ => Self::node(*left, value, *right),
            },
            other => other,
        }
    }

    /// Rebalances a node after insertion.
    /// Handles four cases: left-left, left-right, right-right, right-left.
    fn rebalance(self) -> Self {
        let bf = self.balance_factor();
        if bf > 1 {
            // Left-heavy: check if left child is right-heavy (left-right case)
            match self {
                Avl::Node {
                    ref left,
                    ref value,
                    ref right,
                    ..
                } if left.balance_factor() < 0 => {
                    // Left-right case: rotate left child left, then rotate right
                    let new_left = (**left).clone().rotate_left();
                    Self::node(new_left, value.clone(), (**right).clone()).rotate_right()
                }
                _ => self.rotate_right(),
            }
        } else if bf < -1 {
            // Right-heavy: check if right child is left-heavy (right-left case)
            match self {
                Avl::Node {
                    ref left,
                    ref value,
                    ref right,
                    ..
                } if right.balance_factor() > 0 => {
                    // Right-left case: rotate right child right, then rotate left
                    let new_right = (**right).clone().rotate_right();
                    Self::node((**left).clone(), value.clone(), new_right).rotate_left()
                }
                _ => self.rotate_left(),
            }
        } else {
            self
        }
    }

    /// Inserts a value, returning a new balanced AVL tree.
    /// Duplicates are ignored.
    pub fn insert(&self, x: T) -> Self {
        match self {
            Avl::Empty => Self::node(Avl::Empty, x, Avl::Empty),
            Avl::Node {
                left, value, right, ..
            } => match x.cmp(value) {
                Ordering::Less => {
                    Self::node(left.insert(x), value.clone(), (**right).clone()).rebalance()
                }
                Ordering::Greater => {
                    Self::node((**left).clone(), value.clone(), right.insert(x)).rebalance()
                }
                Ordering::Equal => self.clone(),
            },
        }
    }

    /// In-order traversal returns sorted elements.
    pub fn inorder(&self) -> Vec<T> {
        match self {
            Avl::Empty => vec![],
            Avl::Node {
                left, value, right, ..
            } => {
                let mut result = left.inorder();
                result.push(value.clone());
                result.extend(right.inorder());
                result
            }
        }
    }

    /// Builds an AVL tree from an iterator.
    pub fn build(items: impl IntoIterator<Item = T>) -> Self {
        items
            .into_iter()
            .fold(Avl::empty(), |tree, x| tree.insert(x))
    }

    /// Checks if the tree satisfies the AVL balance invariant.
    pub fn is_balanced(&self) -> bool {
        match self {
            Avl::Empty => true,
            Avl::Node { left, right, .. } => {
                let bf = self.balance_factor();
                (-1..=1).contains(&bf) && left.is_balanced() && right.is_balanced()
            }
        }
    }
}

impl<T: Ord + Clone> Default for Avl<T> {
    fn default() -> Self {
        Self::empty()
    }
}

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

    #[test]
    fn test_empty_tree() {
        let tree: Avl<i32> = Avl::empty();
        assert_eq!(tree.inorder(), Vec::<i32>::new());
        assert_eq!(tree.height(), 0);
        assert!(tree.is_balanced());
    }

    #[test]
    fn test_single_element() {
        let tree = Avl::empty().insert(42);
        assert_eq!(tree.inorder(), vec![42]);
        assert_eq!(tree.height(), 1);
        assert!(tree.is_balanced());
    }

    #[test]
    fn test_sorted_insertion_stays_balanced() {
        // Inserting in sorted order would degrade a plain BST to a linked list.
        // AVL rotations keep it balanced.
        let tree = Avl::build(1..=7);
        assert_eq!(tree.inorder(), vec![1, 2, 3, 4, 5, 6, 7]);
        assert!(tree.is_balanced());
        assert!(tree.height() <= 4); // log2(7) + 1 ≈ 3.8
    }

    #[test]
    fn test_reverse_insertion_stays_balanced() {
        let tree = Avl::build((1..=8).rev());
        assert_eq!(tree.inorder(), vec![1, 2, 3, 4, 5, 6, 7, 8]);
        assert!(tree.is_balanced());
    }

    #[test]
    fn test_ocaml_example() {
        let tree = Avl::build([7, 3, 9, 1, 5, 8, 10, 2]);
        assert_eq!(tree.inorder(), vec![1, 2, 3, 5, 7, 8, 9, 10]);
        assert!(tree.is_balanced());
    }

    #[test]
    fn test_duplicate_insert() {
        let tree = Avl::build([3, 1, 3, 2, 1]);
        assert_eq!(tree.inorder(), vec![1, 2, 3]);
        assert!(tree.is_balanced());
    }
}

type 'a avl = Empty | Node of 'a avl * 'a * 'a avl * int

let height = function Empty -> 0 | Node (_, _, _, h) -> h
let node l v r = Node (l, v, r, 1 + max (height l) (height r))
let balance t = match t with Empty -> 0 | Node (l,_,r,_) -> height l - height r

let rotate_right = function
  | Node (Node (ll, lv, lr, _), v, r, _) -> node (node ll lv lr) v r
  | t -> t

let rotate_left = function
  | Node (l, v, Node (rl, rv, rr, _), _) -> node l v (node rl rv rr)
  | t -> t

let rebalance t = match balance t with
  | b when b > 1 -> rotate_right t
  | b when b < -1 -> rotate_left t
  | _ -> t

let rec insert x = function
  | Empty -> node Empty x Empty
  | Node (l, v, r, _) ->
    if x < v then rebalance (node (insert x l) v r)
    else if x > v then rebalance (node l v (insert x r))
    else node l v r

let rec inorder = function
  | Empty -> [] | Node (l,v,r,_) -> inorder l @ [v] @ inorder r

let () =
  let t = List.fold_left (fun t x -> insert x t) Empty [7;3;9;1;5;8;10;2] in
  assert (inorder t = [1;2;3;5;7;8;9;10]);
  (* Balance check: height should be small *)
  assert (height t <= 4);
  print_endline "ok"

✓ Tests Rust test suite

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

    #[test]
    fn test_empty_tree() {
        let tree: Avl<i32> = Avl::empty();
        assert_eq!(tree.inorder(), Vec::<i32>::new());
        assert_eq!(tree.height(), 0);
        assert!(tree.is_balanced());
    }

    #[test]
    fn test_single_element() {
        let tree = Avl::empty().insert(42);
        assert_eq!(tree.inorder(), vec![42]);
        assert_eq!(tree.height(), 1);
        assert!(tree.is_balanced());
    }

    #[test]
    fn test_sorted_insertion_stays_balanced() {
        // Inserting in sorted order would degrade a plain BST to a linked list.
        // AVL rotations keep it balanced.
        let tree = Avl::build(1..=7);
        assert_eq!(tree.inorder(), vec![1, 2, 3, 4, 5, 6, 7]);
        assert!(tree.is_balanced());
        assert!(tree.height() <= 4); // log2(7) + 1 ≈ 3.8
    }

    #[test]
    fn test_reverse_insertion_stays_balanced() {
        let tree = Avl::build((1..=8).rev());
        assert_eq!(tree.inorder(), vec![1, 2, 3, 4, 5, 6, 7, 8]);
        assert!(tree.is_balanced());
    }

    #[test]
    fn test_ocaml_example() {
        let tree = Avl::build([7, 3, 9, 1, 5, 8, 10, 2]);
        assert_eq!(tree.inorder(), vec![1, 2, 3, 5, 7, 8, 9, 10]);
        assert!(tree.is_balanced());
    }

    #[test]
    fn test_duplicate_insert() {
        let tree = Avl::build([3, 1, 3, 2, 1]);
        assert_eq!(tree.inorder(), vec![1, 2, 3]);
        assert!(tree.is_balanced());
    }
}

Deep Comparison

OCaml vs Rust: AVL Tree — Self-Balancing BST

Side-by-Side Code

OCaml

type 'a avl = Empty | Node of 'a avl * 'a * 'a avl * int

let height = function Empty -> 0 | Node (_, _, _, h) -> h
let node l v r = Node (l, v, r, 1 + max (height l) (height r))

let rotate_right = function
  | Node (Node (ll, lv, lr, _), v, r, _) -> node (node ll lv lr) v r
  | t -> t

let rec insert x = function
  | Empty -> node Empty x Empty
  | Node (l, v, r, _) ->
    if x < v then rebalance (node (insert x l) v r)
    else if x > v then rebalance (node l v (insert x r))
    else node l v r

Rust (idiomatic with named fields)

#[derive(Debug, Clone, PartialEq)]
pub enum Avl<T> {
    Empty,
    Node { left: Box<Avl<T>>, value: T, right: Box<Avl<T>>, height: i32 },
}

impl<T: Ord + Clone> Avl<T> {
    fn rotate_right(self) -> Self {
        match self {
            Avl::Node { left, value, right, .. } => match *left {
                Avl::Node { left: ll, value: lv, right: lr, .. } =>
                    Self::node(Self::node(*ll, lv, *lr), value, *right),
                _ => Self::node(*left, value, *right),
            },
            other => other,
        }
    }

    pub fn insert(&self, x: T) -> Self {
        match self {
            Avl::Empty => Self::node(Avl::Empty, x, Avl::Empty),
            Avl::Node { left, value, right, .. } => match x.cmp(value) {
                Ordering::Less => Self::node(left.insert(x), value.clone(), (**right).clone()).rebalance(),
                Ordering::Greater => Self::node((**left).clone(), value.clone(), right.insert(x)).rebalance(),
                Ordering::Equal => self.clone(),
            },
        }
    }
}

Type Signatures

Concept	OCaml	Rust
Tree type	`'a avl`	`Avl<T>`
Height	`'a avl -> int`	`fn height(&self) -> i32`
Insert	`'a -> 'a avl -> 'a avl`	`fn insert(&self, x: T) -> Self`
Rotate	`'a avl -> 'a avl`	`fn rotate_right(self) -> Self`
Balance factor	`'a avl -> int`	`fn balance_factor(&self) -> i32`

Key Insights

Rotation = ownership transfer: OCaml creates new nodes from pattern-matched pieces; Rust moves ownership through self by value, making the restructuring zero-copy where possible

Two-level pattern matching: OCaml matches Node(Node(ll,lv,lr,_), v, r, _) in one arm; Rust needs nested match since it can't destructure two levels of Box at once

Named vs positional fields: OCaml's Node of 'a avl * 'a * 'a avl * int uses positional access; Rust's named struct fields (left, value, right, height) are self-documenting

Balance invariant verification: Both languages can express is_balanced as a recursive check; Rust's type system doesn't encode the invariant statically (both rely on runtime checks)

Clone cost transparency: Rust's .clone() on unchanged subtrees makes the persistence cost visible; OCaml shares references silently via GC

When to Use Each Style

Use AVL tree when: you need guaranteed O(log n) operations and care about worst-case performance Use standard BTreeMap when: you want a production-ready balanced tree — Rust's stdlib provides one

Exercises

Add a rank method that returns how many elements in the tree are strictly less than a given value, by augmenting each node with its subtree count.

Implement split that divides the AVL tree into two balanced trees: one with all elements less than a pivot, and one with all elements greater.

Write a property-based test suite that verifies: (1) BST ordering, (2) balance factor ≤ 1 at every node, and (3) to_vec returns elements in sorted order — all after a random sequence of inserts and deletes.

Open Source Repos

functional-rust

View the source for this example on GitHub — OCaml and Rust side by side in the repo.

Rust