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)?

/**
 * 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> res = new ArrayList<Integer>();
        PriorityQueue<Integer> queue = new PriorityQueue<>(
            new Comparator<Integer>() {
                public int compare(Integer i1, Integer i2) {
                    double sign = Math.abs(i1 - target) - Math.abs(i2 - target);
                    if (sign > 0) return 1;
                    else if (sign < 0) return -1;
                    else return 0;
                }
            }
            );
        visit(root, queue);
        while (k-- > 0) {
            res.add(queue.poll());
        }
        return res;
    }
 
    private void visit(TreeNode node, PriorityQueue<Integer> queue) {
        if (node == null) return;
        queue.add(node.val);
        visit(node.left, queue);
        visit(node.right, queue);
    }
}

======

/**
 * 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> res = new ArrayList<Integer>();
        PriorityQueue<Integer> queue = new PriorityQueue<>(
            new Comparator<Integer>() {
                public int compare(Integer i1, Integer i2) {
                    double sign = Math.abs(i1 - target) - Math.abs(i2 - target);
                    if (sign > 0) return -1;
                    else if (sign < 0) return 1;
                    else return 0;
                }
            }
            );
        visit(root, queue, k, target);
        while (k-- > 0) {
            res.add(0, queue.poll());
        }
        return res;
    }
 
    private void visit(TreeNode node, PriorityQueue<Integer> queue, int k, double target) {
        if (node == null) return;
        if (queue.size() == k) {
            if (Math.abs(node.val - target) < Math.abs(queue.peek() - target)) {
                queue.poll();
                queue.offer(node.val);
            }
        } else {
            queue.offer(node.val);
        }
        visit(node.left, queue, k, target);
        visit(node.right, queue, k, target);
    }
}

======

/**
 * 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> res = new ArrayList<>();
        List<Integer> small = new ArrayList<>();
        List<Integer> big = new ArrayList<>();
       
        visitSmall(root, target, small);
        visitBig(root, target, big);
       
        int p1 = small.size() - 1;
        int p2 = big.size() - 1;
        while (p1 >= 0 || p2 >= 0) {
            if (p1 < 0 || (p2 >= 0 && Math.abs(small.get(p1) - target) > Math.abs(big.get(p2) - target))) {
                res.add(big.get(p2));
                p2--;
            } else {
                res.add(small.get(p1));
                p1--;
            }
            if (res.size() == k) break;
        }
        return res;
    }
   
    private void visitSmall(TreeNode node, double target, List<Integer> res) {
        if (node == null) return;
        visitSmall(node.left, target, res);
        if (node.val > target) {
            return;
        }
        res.add(node.val);
        visitSmall(node.right, target, res);
    }
   
    private void visitBig(TreeNode node, double target, List<Integer> res) {
        if (node == null) return;
        visitBig(node.right, target, res);
        if (node.val <= target) {
            return;
        }
        res.add(node.val);
        visitBig(node.left, target, res);
    }
}