[LintCode] 901. 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.

Example

Example 1:

Input:
{1}
0.000000
1
Output:
[1]
Explanation：
Binary tree {1},  denote the following structure:
 1

Example 2:

Input:
{3,1,4,#,2}
0.275000
2
Output:
[1,2]
Explanation：
Binary tree {3,1,4,#,2},  denote the following structure:
  3
 /  
1    4
 
  2

Challenge

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

Notice

• 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.

/**
 * 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> nodeList = new ArrayList<>();
        dfs(root, nodeList);
        int i = 0;
        while (i < nodeList.size()) {
            if (nodeList.get(i) >= target) {
                break;
            }
            i+= 1;
        }
        int left = i - 1, right = i;
        List<Integer> res = new ArrayList<>();
        while (k > 0) {
            if (left >= 0 && (right >= nodeList.size() || target - nodeList.get(left) < nodeList.get(right) - target)) {
                res.add(nodeList.get(left--));
            } else {
                res.add(nodeList.get(right++));
            }
            k -= 1;
        }
        return res;
    }
    
    private void dfs(TreeNode root, List<Integer> list) {
        if (root == null) {
            return;
        }
        dfs(root.left, list);
        list.add(root.val);
        dfs(root.right, list);
    }
}