Given a binary tree, create a list of all the nodes at each depth — Recursive version

Another FAQ on the programming interview.

Question

Given a binary search tree, design an algorithm which creates a linked list of all the nodes at each depth (eg, if you have a tree with depth D, you’ll have D linked lists).

Solution

Cracking the Coding Interview (4ed) (see page 126) provided a non-recursive algorithm. For me, a recursive version is neater and much easier to understand.

public List<List<TreeNode>> findLevelLinkList(TreeNode root) {
  return findLevelLinkListHelper(
      root,
      0,
      new ArrayList<List<TreeNode>>());
}

public List<List<TreeNode>> findLevelLinkListHelper(
    TreeNode root,
    int level,
    List<List<TreeNode>> lists) {
  if (root != null) {
    if (level >= lists.size()) {
      lists.add(new LinkedList<TreeNode>());
    }
    lists.get(level).add(root);
    findLevelLinkListHelper(root.left,  level + 1, lists);
    findLevelLinkListHelper(root.right, level + 1, lists);
  }
  return lists;
}