Counting nodes in a tree in Java

Question

First of all, I swear this is not homework, it's a question I was asked in an interview. I think I made a mess of it (though I did realise the solution requires recursion). Here is the question:

Implement the count() method which returns the number of nodes in a tree. If a node doesn't have either a left or right child, the relevant getXXChild() method will return null

class Tree {

  Tree getRightChild() {
    // Assume this is already implemented
  }

  Tree getLeftChild() {
    // Assume this is already implemented
  }

  int count() {
    // Implement me
  }
}

My reason for asking the question is simply curious to see the correct solution, and thereby measure how bad mine was.

Cheers, Tony

Answer 1

Of course, if you want to avoid visiting every node in your tree when you count, and processing time is worth more to you than memory, you can cheat by creating your counts as you build your tree.

1. Have an int count in each node, initialized to one, which respresents the number of nodes in the subtree rooted in that node.

2. When you insert a node, before returning from your recursive insert routine, increment the count at the current node.

i.e.

public void insert(Node root, Node newNode) {
  if (newNode.compareTo(root) > 1) {
    if (root.right != null) 
      insert(root.right, newNode);
    else
      root.right = newNode;
  } else {
    if (root.left != null)
      insert(root.left, newNode);
    else
      root.left = newNode;
  }
  root.count++;
}

Then getting the count from any point just involves a lookup of node.count

Answer 2

int count() {
  Tree right = getRightChild();
  Tree left = getLeftChild();
  int c = 1;                                      // count yourself!
  if ( right != null ) c += right.count();        // count sub trees
  if ( left != null ) c += left.count();          // ..
  return c;
}

Answer 3

A trivial recursive solution:

int count() {
   Tree l = getLeftTree();
   Tree r = getRightTree();
   return 1 + (l != null ? l.count() : 0) + (r != null ? r.count() : 0);
}

A less trivial non-recursive one:

int count() {
    Stack<Tree> s = new Stack<Tree>();
    s.push(this);
    int cnt = 0;
    while (!s.empty()) {
        Tree t = s.pop();
        cnt++;
        Tree ch = getLeftTree();
        if (ch != null) s.push(ch); 
        ch = getRightTree();
        if (ch != null) s.push(ch); 
    }
    return cnt;
}

The latter is probably slightly more memory-efficient, because it replaces recursion with a stack and an iteration. It's also probably faster, but its hard to tell without measurements. A key difference is that the recursive solution uses the stack, while the non-recursive solution uses the heap to store the nodes.

Edit: Here's a variant of the iterative solution, which uses the stack less heavily:

int count() {
    Tree t = this;
    Stack<Tree> s = new Stack<Tree>();
    int cnt = 0;
    do {
        cnt++;
        Tree l = t.getLeftTree();
        Tree r = t.getRightTree();
        if (l != null) {
            t = l;
            if (r != null) s.push(r);
        } else if (r != null) {
            t = r;
        } else {
            t = s.empty() ? null : s.pop();
        }
    } while (t != null);
    return cnt;
}

Whether you need a more efficient or a more elegant solution naturally depends on the size of your trees and on how often you intend to use this routine. Rembemer what Hoare said: "premature optimization is the root of all evil."