Writing a BST in Java

Last week you worked through the BST on paper. This week we put the same algorithms into Java. The exam has two kinds of questions: the output of given code (traversals after insert/remove calls) and “write a method”: short recursive methods such as size, height, count of elements in a range, k-th smallest, BST check. Someone who knows the patterns on this page can write a new method in 5 minutes too.

Three golden rules:

1. A recursive method returns a node, and the caller assigns it: t.left = insert(x, t.left);. Java passes references by value.
2. Every method deals with the t == null case first. If the base case is wrong, everything else collapses.
3. If you can prune one side of a BST, prune it. Operations like search, range and floor are O(h); operations that traverse the whole tree are O(n).

The skeleton of the class

The key type is generic. Comparison is done with compareTo, not <, so the type must be Comparable.

public class BinarySearchTree<T extends Comparable<? super T>> {

    private static class BinaryNode<T> {
        T element;
        BinaryNode<T> left, right;
        BinaryNode(T e) { element = e; }
    }

    private BinaryNode<T> root;   // null in an empty tree

    public boolean isEmpty() { return root == null; }
    public void makeEmpty()  { root = null; }
}

• BinaryNode is a private static nested class: it is invisible from outside and is not tied to a tree object.
• <? super T>: it means accept it even if not T itself but a superclass of T is Comparable. Just memorize it.
• makeEmpty only sets the root to null. The garbage collector cleans up the nodes.

Public driver, private helper

Each operation consists of two methods. The public one is open to the user and passes only the root. The private one does the real recursion and takes the root of a subtree as a parameter.

public void insert(T x)      { root = insert(x, root); }
public void remove(T x)      { root = remove(x, root); }
public boolean contains(T x) { return contains(x, root); }

The user never sees the concept of a node. In operations where the root can change (the first insertion, deleting the root), the root = ... assignment is required.

contains, findMin, findMax

private boolean contains(T x, BinaryNode<T> t) {
    if (t == null) return false;
    int c = x.compareTo(t.element);
    if (c < 0) return contains(x, t.left);
    if (c > 0) return contains(x, t.right);
    return true;                         // c == 0: found
}

private BinaryNode<T> findMin(BinaryNode<T> t) {
    if (t.left == null) return t;        // recursive
    return findMin(t.left);
}

private BinaryNode<T> findMax(BinaryNode<T> t) {
    while (t.right != null) t = t.right; // with a loop
    return t;
}

compareTo returns negative, zero or positive. Do not check with == -1: the contract guarantees only the sign. Since contains is tail recursion, it is easy to turn into a loop. The private versions of findMin and findMax return a node; the public versions throw an exception on an empty tree and return element.

insert

private BinaryNode<T> insert(T x, BinaryNode<T> t) {
    if (t == null) return new BinaryNode<>(x);   // new leaf
    int c = x.compareTo(t.element);
    if (c < 0)      t.left  = insert(x, t.left);
    else if (c > 0) t.right = insert(x, t.right);
    // c == 0: duplicate, do nothing
    return t;
}

The logic step by step: the call insert(25, root) goes from 40 to 20, and from 20 to 30. The left of 30 is null, insert(25, null) returns a new node, and the assignment 30.left = new node is made. On the way back up, the other assignments (20.right = 30, 40.left = 20) rewrite links that already exist, which breaks nothing.

PITFALL · recursion without assignment. If you write insert(x, t.left); and do not assign the returned value, the new node is left dangling. Doing t = new BinaryNode<>(x) inside the method does not change the caller’s reference. Result: the tree has only the root.

remove

private BinaryNode<T> remove(T x, BinaryNode<T> t) {
    if (t == null) return null;                  // absent: do nothing
    int c = x.compareTo(t.element);
    if (c < 0)      t.left  = remove(x, t.left);
    else if (c > 0) t.right = remove(x, t.right);
    else if (t.left != null && t.right != null) {        // two children
        t.element = findMin(t.right).element;
        t.right = remove(t.element, t.right);
    } else {                                     // 0 or 1 child
        t = (t.left != null) ? t.left : t.right;
    }
    return t;
}

The ternary expression in the last line handles two cases at once. At a leaf both children are null, so t becomes null and the parent’s reference is set to null. With one child, t becomes that child and the parent attaches directly to it.

In the two-children case a second search is done in the right subtree. The total is still $O(h)$: once down to find the successor, and once to delete it.

Traversals, size, height

private void inorder(BinaryNode<T> t, StringBuilder sb) {
    if (t == null) return;
    inorder(t.left, sb);
    sb.append(t.element).append(' ');
    inorder(t.right, sb);
}

private int size(BinaryNode<T> t) {
    return t == null ? 0 : 1 + size(t.left) + size(t.right);
}

private int height(BinaryNode<T> t) {
    return t == null ? -1 : 1 + Math.max(height(t.left), height(t.right));
}

size is $O(n)$ because it traverses the whole tree on every call. If it is called often, adding an int size field to the class and updating it in insert and remove gives $O(1)$. But then you need to know whether insert really added something (or was a duplicate).

A worked example

BinarySearchTree<Integer> bst = new BinarySearchTree<>();
int[] keys = {40, 20, 60, 10, 30, 50, 70, 25};
for (int k : keys) bst.insert(k);
System.out.println("inorder  : " + bst.inorder());
System.out.println("preorder : " + bst.preorder());
System.out.println("postorder: " + bst.postorder());
System.out.println("size=" + bst.size() + " height=" + bst.height()
        + " min=" + bst.findMin() + " max=" + bst.findMax());
System.out.println("contains 25? " + bst.contains(25)
        + ", contains 35? " + bst.contains(35));
bst.remove(20);
System.out.println("after remove(20): " + bst.preorder());

inorder  : 10 20 25 30 40 50 60 70
preorder : 40 20 10 30 25 60 50 70
postorder: 10 25 30 20 50 70 60 40
size=8 height=3 min=10 max=70
contains 25? true, contains 35? false
after remove(20): 40 25 10 30 60 50 70

  after insertion      after remove(20)
       40                    40
     /    \                /    \
   20      60            25      60
  /  \    /  \          /  \    /  \
10   30  50   70      10   30  50   70
     /
   25

The successor of 20 is the smallest of the right subtree, which is 25. The same class works with String too: inserting "mango kiwi apple peach banana" gives the inorder apple banana kiwi mango peach, because String’s compareTo is alphabetical.

Write-a-method patterns

Most of the methods asked in exams fit one of three patterns.

The sample tree used in the questions on this page (the keys 50, 30, 70, 20, 40, 60, 80, 35, 45, 65 were inserted in this order):

            50
          /    \
        30      70
       /  \    /  \
     20   40  60   80
         / \    \
       35  45    65

TIP · think of the empty tree and the single node first. Run the method you wrote in your head on null and on a single leaf. If these two cases are right, the recursion usually handles the rest.

Common mistakes

BinarySearchTree<Integer> t = new BinarySearchTree<>();
for (int x : new int[]{33, 18, 47, 9, 25, 40, 52, 18}) t.insert(x);
System.out.println(t.findMin() + " " + t.findMax() + " " + t.size());

BinarySearchTree<Integer> u = new BinarySearchTree<>();
for (int x : new int[]{33, 18, 47, 9, 25, 40, 52, 21}) u.insert(x);
u.remove(18);
u.remove(47);
System.out.println(u.preorder());
System.out.println(u.height());

static int mystery(Node t) {
    if (t == null) return 0;
    int c = ((t.left == null) != (t.right == null)) ? 1 : 0;
    return c + mystery(t.left) + mystery(t.right);
}

static Node badInsert(Node t, int x) {
    if (t == null) return new Node(x);
    if (x < t.element) badInsert(t.left, x);
    else if (x > t.element) badInsert(t.right, x);
    return t;
}

static int size(Node t) {
    if (t == null) return 0;
    return 1 + size(t.left) + size(t.right);
}

static int countInRange(Node t, int lo, int hi) {
    if (t == null) return 0;
    if (t.element < lo) return countInRange(t.right, lo, hi);
    if (t.element > hi) return countInRange(t.left, lo, hi);
    return 1 + countInRange(t.left, lo, hi)
             + countInRange(t.right, lo, hi);
}

static boolean isBST(Node t) {
    return isBST(t, Long.MIN_VALUE, Long.MAX_VALUE);
}
static boolean isBST(Node t, long lo, long hi) {
    if (t == null) return true;
    if (t.element <= lo || t.element >= hi) return false;
    return isBST(t.left, lo, t.element)
        && isBST(t.right, t.element, hi);
}

static int count, answer;
static int kthSmallest(Node t, int k) {
    count = 0; answer = -1;
    kth(t, k);
    return answer;
}
static void kth(Node t, int k) {
    if (t == null || count >= k) return;
    kth(t.left, k);
    if (++count == k) { answer = t.element; return; }
    kth(t.right, k);
}

static Integer floor(Node t, int x) {
    Integer best = null;
    while (t != null) {
        if (t.element == x) return x;
        if (t.element < x) { best = t.element; t = t.right; }
        else t = t.left;
    }
    return best;
}

static Node lca(Node t, int a, int b) {
    while (t != null) {
        if (a < t.element && b < t.element) t = t.left;
        else if (a > t.element && b > t.element) t = t.right;
        else return t;
    }
    return null;
}