Friday, November 7, 2014

Facebook: Print n-th node in a binary tree

Write a function that takes 2 arguments: a binary tree and an integers, it should return the n-th element in the inorder traversal of the binary tree.

Understand the problem:
The problem asks for returning the n-th element in a binary tree. n starts from 1. 

Recursive Solution:
public class Solution {
    private int count = 0;
    private int num = 0;
    public int nthElement(TreeNode root, int n) {
        if (n <= 0) {
            return Integer.MIN_VALUE;
        }
        
        if (nthElementHelper(root, n) == false) {
            return Integer.MIN_VALUE; // not found
        } else {
            return num;
        }
    }
    
    private boolean nthElementHelper(TreeNode root, int n) {
        if (root == null) {
            return false;
        }
        
        if (nthElementHelper(root.left, n) == true) {
            return true;
        }
        
        count++;
        if (count == n) {
            num = root.val;
            return true;
        }
        
        if (nthElementHelper(root.right, n) == true) {
            return true;
        }
        
        return false;
    }

    public static void main(String[] args) {
      Solution sol = new Solution();
      TreeNode root = new TreeNode(1);
      root.left = new TreeNode(2);
      root.right = new TreeNode(3);
      root.left.left = new TreeNode(4);
      root.left.right = new TreeNode(5);
      root.right.right = new TreeNode(6);

      int result = sol.nthElement(root, 6);
      
      System.out.println("Result is " + result);
      
    }
}

Discussion:
The trick in the solution is to stop traversing the tree whenever we found the nth node. How do we do it? We use returning flag as true or false to mark. If we found return true, and stop the process, else return false and continue find it. The advantage of using this is we can know if we eventually found that number or not, by checking the return flag. Therefore, we don't need to assume that n is less or equal to the number of nodes in the tree.

Iterative Solution:
public class Solution {
    public int nthElement(TreeNode root, int n) {
        if (n <= 0 || root == null) {
            return Integer.MIN_VALUE;
        }
        
        Stack<TreeNode> stack = new Stack<TreeNode>();
        TreeNode p = root;
        int count = 0;
        int result = Integer.MIN_VALUE;
         
        while (p != null || !stack.isEmpty()) {
            if (p != null) {
                stack.push(p);
                p = p.left;
            } else {
                TreeNode curr = stack.pop();
                count++;
                if (count == n) {
                    result = curr.val;
                    return result;
                }
                p = curr.right;
            }
        }
        
        return result;
    }
}

No comments:

Post a Comment