Sunday, November 16, 2014

Facebook: Implement a iterator of a BST

/**
 * Definition of TreeNode:
 * public class TreeNode {
 *     public int val;
 *     public TreeNode left, right;
 *     public TreeNode(int val) {
 *         this.val = val;
 *         this.left = this.right = null;
 *     }
 * }
 * Example of iterate a tree:
 * Solution iterator = new Solution(root);
 * while (iterator.hasNext()) {
 *    TreeNode node = iterator.next();
 *    do something for node
 * } 
 */
import java.util.*;
public class Solution {
    private Stack<TreeNode> stack;
    private TreeNode curt;
    
    private void moveNext() {
        if (curt.right == null) {
            while (!stack.empty()) {
                TreeNode peek  = stack.peek();
                if (peek.left == curt) {
                    curt = stack.pop();
                    return;
                }
                if (peek.right == curt) {
                    curt = stack.pop();
                }
            }
            if (stack.empty()) {
                curt = null;
                return;
            }
        } else {
            stack.push(curt);
            curt = curt.right;
            while (curt.left != null) {
                stack.push(curt);
                curt = curt.left;
            }
        }
    }
    
    // @param root: The root of binary tree.
    public Solution(TreeNode root) {
        stack = new Stack<TreeNode>();
        curt = new TreeNode(0);
        curt.right = root;
        moveNext();
    }

    //@return: True if there has next node, or false
    public boolean hasNext() {
        return curt != null;
    }
    
    //@return: return next node
    public TreeNode next() {
        TreeNode temp = curt;
        moveNext();
        return temp;
    }

    public static void main(String[] args) {
        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);

        Solution iterator = new Solution(root);
        
        while (iterator.hasNext()) {
            TreeNode cur = iterator.next();
            System.out.println(cur.val);
        }        
    }
}

No comments:

Post a Comment