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

Solution 1

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;
}

Solution 2

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."

Solution 3

I like this better because it reads:

return count for left + count for rigth + 1

  int count() {
      return  countFor( getLeftChild() ) + countFor( getRightChild() ) + 1;
  }
  private int countFor( Tree tree )  { 
       return tree == null ? 0 : tree.count();
  }

A little more towards literate programming.

BTW, I don't like the getter/setter convention that is so commonly used on Java, I think a using leftChild() instead would be better:

  return countFor( leftChild() ) + countFor( rightChild() ) + 1;

Just like Hoshua Bloch explains here http://www.youtube.com/watch?v=aAb7hSCtvGw at min. 32:03

If you get it rigth your code reads...

BUT, I have to admit the get/set convention is now almost part of the language. :)

For many other parts, following this strategy creates self documenting code, which is something good.

Tony: I wonder, what was your answer in the interview.

Solution 4

class Tree {

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

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

  int count() {
   return 1 
      + getRightChild() == null? 0 : getRightChild().count()
      + getLeftChild() == null? 0 : getLeftChild().count();
  }
}

Solution 5

Something like this should work:

int count()
{
    int left = getLeftChild() == null ? 0 : getLeftChild().count();
    int right = getRightChild() == null ? 0 : getRightCHild().count();

    return left + right + 1;
}

View more solutions

Comments

[Tim Frey]

This won't compile with Java. Unfortunately, you need to explicitly test if something == null, so you need l== null and r == null here.

[Tim Frey]

Pfft what are you guys talking about, I always added 1 ;-) Thanks!

[Bill Karwin]

Edited "return count" to "return c". :-)

[Tim Frey]

Your if statements need to explicitly check == null, otherwise this won't work for Java.

[mannicken]

Yep, recalled that when I posted :)

[David Z]

I wouldn't write this as a one-liner in practice but the fact that you can makes a point about the simplicity... +1

[Steven Robbins]

Why not? Can't beat a bit of ternary abuse :-)

[David Hanak]

Thanks! It was some time ago I actively used Java.

[OscarRyz]

Who said count should return null? The problem says "The relevant method" so it is referring to the "leftNode/rightNode". That was your point or a missed something?

[Tim Frey]

How about allowing nulls to be pushed onto the stack, but only incrementing the counter if the object popped off is not null. You can remove 1 if statement this way, and also you can then do s.push(getLeft()) and s.push(getRight())

[Kip]

I too thought he meant the count() method should return null if there are no child nodes. It wasn't clear until the second reading. When he said "relevant method", I assumed he meant the method relevant to this question.

[Kip]

but that's as efficient as it gets, without caching higher-level counts, right?

[Kip]

If you wrote that during an interview, i would be annoyed not impressed.

[Tim Frey]

The algorithm is fine, it just doesn't really answer the original question. Plus, haveLeft and haveRight, while obvious, aren't declared here.

[Steven Robbins]

@Kip Not as annoyed as I would be if I ended up in a Java interview :-)

[oxbow_lakes]

-1 points on going against the JavaBean idiom, which has proven hugely useful. (I upscored you anyway!)

[OscarRyz]

@oxbow: Yeap, the Java bean model never got the desired effect ( which was to have truly standalone beans ) but the idiom was been very beneficial, specially in webdevelopment.

[Platinum Azure]

You have to do more processing on the result of each of the count calls before returning (adding them to the return register's current value), so I don't think you have a proper tail call here. :-(

[richj]

Platinum Azure: Thanks for the comment. It looks like a fatal flaw in my idea. If TCO can't deal with the accumulation in the return value then I guess I need a better optimization :-(

[QuantumMechanic]

I thought it was Knuth who said that! :)