Leetcode: 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.

Brute-force solution:
The straight-forward solution would be to use a heap. We just treat the BST just as a usual array and do a in-order traverse. Then we compare the current element with the minimum element in the heap, the same as top k problem.

Code (Java):

/**
 * Definition for a binary tree node.
 * public class TreeNode {
 *     int val;
 *     TreeNode left;
 *     TreeNode right;
 *     TreeNode(int x) { val = x; }
 * }
 */
public class Solution {
    private PriorityQueue<Integer> minPQ;
    private int count = 0;
    public List<Integer> closestKValues(TreeNode root, double target, int k) {
        minPQ = new PriorityQueue<Integer>(k);
        List<Integer> result = new ArrayList<Integer>();
        
        inorderTraverse(root, target, k);
        
        // Dump the pq into result list
        for (Integer elem : minPQ) {
            result.add(elem);
        }
        
        return result;
    }
    
    private void inorderTraverse(TreeNode root, double target, int k) {
        if (root == null) {
            return;
        }
        
        inorderTraverse(root.left, target, k);
        
        if (count < k) {
            minPQ.offer(root.val);
        } else {
            if (Math.abs((double) root.val - target) < Math.abs((double) minPQ.peek() - target)) {
                minPQ.poll();
                minPQ.offer(root.val);
            }
        }
        count++;
        
        inorderTraverse(root.right, target, k);
    }
}

Analysis:
The time complexity would be O(k + (n - k) logk). 
Space complexity is O(k).

A time linear solution:

/**
 * 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) {
        List<Integer> result = new ArrayList<>();
        if (root == null) {
            return result;
        }
        
        Stack<Integer> precedessor = new Stack<>();
        Stack<Integer> successor = new Stack<>();
        
        getPredecessor(root, target, precedessor);
        getSuccessor(root, target, successor);
        
        for (int i = 0; i < k; i++) {
            if (precedessor.isEmpty()) {
                result.add(successor.pop());
            } else if (successor.isEmpty()) {
                result.add(precedessor.pop());
            } else if (Math.abs((double) precedessor.peek() - target) < Math.abs((double) successor.peek() - target)) {
                result.add(precedessor.pop());
            } else {
                result.add(successor.pop());
            }
        }
        
        return result;
    }
    
    private void getPredecessor(TreeNode root, double target, Stack<Integer> precedessor) {
        if (root == null) {
            return;
        }
        
        getPredecessor(root.left, target, precedessor);
        
        if (root.val > target) {
            return;
        }
        
        precedessor.push(root.val);
        
        getPredecessor(root.right, target, precedessor);
    }
    
    private void getSuccessor(TreeNode root, double target, Stack<Integer> successor) {
        if (root == null) {
            return;
        }
        
        getSuccessor(root.right, target, successor);
        
        if (root.val <= target) {
            return;
        }
        
        successor.push(root.val);
        
        getSuccessor(root.left, target, successor);
    }
}

Update on 4/30/19: Time complexity O(k + logn)

/**
 * 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;
 *     }
 * }
 */

public class Solution {
    /**
     * @param root: the given BST
     * @param target: the given target
     * @param k: the given k
     * @return: k values in the BST that are closest to the target
     */
    public List<Integer> closestKValues(TreeNode root, double target, int k) {
        // write your code here
        List<Integer> ans = new ArrayList<>();
        if (root == null) {
            return ans;
        }

        // step 1: find the closet value and save the path
        Stack<TreeNode> lowerStack = new Stack<>();
        Stack<TreeNode> upperStack = new Stack<>();

        TreeNode p = root;
        while (p != null) {
            if (p.val < target) {
                lowerStack.push(p);
                p = p.right;
            } else {
                upperStack.push(p);
                p = p.left;
            }
        }

        for (int i = 0; i < k; i++) {
            if (lowerStack.isEmpty()) {
                TreeNode top = upperStack.pop();
                ans.add(top.val);
                goUpperNext(top.right, upperStack);
            
            } else if (upperStack.isEmpty()) {
                TreeNode top = lowerStack.pop();
                ans.add(top.val);
                goLowerNext(top.left, lowerStack);
            } else if (upperStack.peek().val - target <= target - lowerStack.peek().val) {
                TreeNode top = upperStack.pop();
                ans.add(top.val);
                goUpperNext(top.right, upperStack);
            } else if (upperStack.isEmpty() || target - lowerStack.peek().val < upperStack.peek().val - target) {
                TreeNode top = lowerStack.pop();
                ans.add(top.val);
                goLowerNext(top.left, lowerStack);
            }
        }

        return ans;
    }

    private void goUpperNext(TreeNode node, Stack<TreeNode> stack) {
        TreeNode p = node;
        while (p != null) {
            stack.push(p);
            p = p.left;
        }
    }

    private void goLowerNext(TreeNode node, Stack<TreeNode> stack) {
        TreeNode p = node;
        while (p != null) {
            stack.push(p);
            p = p.right;
        }
    }
}