272. Closest Binary Search Tree Value II

Given a non-empty binary search tree and a target value, find k values in the BST that are closest to the target.

Note:

• Given target value is a floating point.
• You may assume k is always valid, that is: k ≤ total nodes.
• You are guaranteed to have only one unique set of k values in the BST that are closest to the target.

Follow up:
Assume that the BST is balanced, could you solve it in less than O(n) runtime (where n = total nodes)?

Hint:

1. Consider implement these two helper functions:
   1. getPredecessor(N), which returns the next smaller node to N.
   2. getSuccessor(N), which returns the next larger node to N.
2. Try to assume that each node has a parent pointer, it makes the problem much easier.
3. Without parent pointer we just need to keep track of the path from the root to the current node using a stack.
4. You would need two stacks to track the path in finding predecessor and successor node separately.

Solution:

Get Next:

if the node has right subtree, the next would be the left-most node of the right sub-tree.

if the node does not have right subtree, go to the parent(stack.peek()), 

if the parent's value bigger than the nodes value, means the node is the left-child, the parent is the next node,

if the parent's value smaller than the nodes value, means the node is the right- child, go to parent's parent util when find the node is in the left subtree of an ancestor.

Get Prev is similar to get Next.

Follow the hint could get the Code:

/**
 * Definition for a binary tree node.
 * public class TreeNode {
 *     int val;
 *     TreeNode left;
 *     TreeNode right;
 *     TreeNode(int x) { val = x; }
 * }
 */
public class Solution {
    public List<Integer> closestKValues(TreeNode root, double target, int k) {
        Stack<TreeNode> pre = new Stack<>();
        Stack<TreeNode> post = new Stack<>();
        TreeNode cur = root;
        double preDiff = 1.0 * Integer.MAX_VALUE;
        TreeNode preNode = null;
        double postDiff = 1.0 * Integer.MAX_VALUE;
        TreeNode postNode = null;
        
        // prepare stacks
        while(cur!= null){
            pre.push(cur);
            post.push(cur);
            if(cur.val > target){
                if(postNode == null || cur.val - target < postDiff){
                    postDiff = cur.val - target;
                    postNode = cur;
                }
                cur = cur.left;
            } else {
                if(preNode == null || target - cur.val < preDiff){
                    preDiff = target - cur.val;
                    preNode = cur;
                }
                cur = cur.right;
            }
        }
        
        // get the Node that is first one smaller than target
        while(!pre.empty() && pre.peek() != preNode){
            pre.pop();
        }
        // get the Node that is the first one bigger than target
        while(!post.empty() && post.peek() != postNode){
            post.pop();
        }
        
        
        
        //get result
        List<Integer> ret = new LinkedList<>();
        
        Integer bigger = getNext(post);
        Integer smaller = getPrev(pre);

        for(int i = 0; i < k; i++){
            if(smaller == null){
                ret.add(bigger);
                bigger = getNext(post);
            } else if(bigger == null){
                ret.add(smaller);
                smaller = getPrev(pre);
            } else {
                if( bigger - target < target - smaller){
                    ret.add(bigger);
                    bigger = getNext(post);
                } else {
                    ret.add(smaller);
                    smaller = getPrev(pre);
                }
            }
        }
        return ret;
    }
    
    //get next, and prepare for the next next
    Integer getNext(Stack<TreeNode> post){
        if(post.empty()) return null;
        TreeNode cur = post.pop();
        int ret = cur.val;
        if(cur.right != null){
            pushAllLeft(post,cur.right);
        } else {
            while(!post.empty() && post.peek().val < ret){
                post.pop();
            }
        }
        return ret;
    }
    
    
    void pushAllLeft(Stack<TreeNode> post, TreeNode root){
        while(root != null){
            post.push(root);
            root = root.left;
        }
    }
    
    //get prev, and prepare for the prev prev
    Integer getPrev(Stack<TreeNode> pre){
        if(pre.empty()) return null;
        TreeNode cur = pre.pop();
        int ret = cur.val;
        if(cur.left != null){
            pushAllRight(pre,cur.left);
        } else {
            while(!pre.empty() && pre.peek().val > ret){
                pre.pop();
            }
        }
        return ret;
    }
    
    void pushAllRight(Stack<TreeNode> pre, TreeNode root){
        while(root != null){
            pre.push(root);
            root = root.right;
        }
    }
}