Monday, November 2, 2015

Zenefits: 2Sum in BST

Input: A binary search tree, consisting of integers. And a number k
Output: True if there are two nodes a and b in the tree, such that a.val + b.val = k. 
Return false otherwise. 

A O(n) time, O(n) space solution:
The most straight-forward solution is first do a in-order traverse of the BST, store the results into an array. Then the problem can be transformed into a 2 sum problem. 
So the time and space complexity for this solution is O(n). 

A O(n) time, O(log n) space solution:
import java.util.*;

public class Solution {
    public boolean twoSumBST(TreeNode root, int k) {
        if (root == null) {
            return false;
        }
        
        Stack<TreeNode> leftStack = new Stack<>();
        Stack<TreeNode> rightStack = new Stack<>();
        
        TreeNode p = root;
        while (p != null) {
            leftStack.push(p);
            p = p.left;
        }
        
        p = root;
        while (p != null) {
            rightStack.push(p);
            p = p.right;
        }
        
        p = leftStack.peek();
        TreeNode q = rightStack.peek();
        
        while (p.val < q.val && p != q) {
            if (p.val + q.val == k) {
                return true;
            } else if (p.val + q.val < k) {
                leftStack.pop();
                if (p.right != null) {
                    p = p.right;
                    while (p != null) {
                        leftStack.push(p);
                        p = p.left;
                    }
                }
                
                if (leftStack.isEmpty()) {
                    return false;
                }
                
                p = leftStack.peek();
            } else if (p.val + q.val > k) {
                rightStack.pop();
                if (q.left != null) {
                    q = q.left;
                    while (q != null) {
                        rightStack.push(q);
                        q = q.right;
                    }
                }
                
                if (rightStack.isEmpty()) {
                    return false;
                }
                
                q = rightStack.peek();
            }
        }
        
        return false;
    }
    
    public static void main(String[] args) {
        Solution solution = new Solution();
        
        TreeNode root = new TreeNode(4);
        root.left = new TreeNode(2);
        root.right = new TreeNode(6);
        
        root.left.left = new TreeNode(1);
        root.left.right = new TreeNode(3);
        
        root.right.left = new TreeNode(5);
        root.right.right = new TreeNode(8);
        
        boolean result = solution.twoSumBST(root, 3);
        
        System.out.println(result);
        
    }
}

class TreeNode {
    int val;
    TreeNode left, right;
        
    public TreeNode(int val) {
        this.val = val;
        left = null;
        right = null;
    }
}

No comments:

Post a Comment