LeetCode - Closest Binary Search Tree Value II

ClosestBinarySearchTreeValueII.java

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

public List<Integer> closestKValues(TreeNode root, double target, int k) {
    List<Integer> res = new ArrayList<Integer>();
    
    Stack<TreeNode> successors = new Stack<TreeNode>(); 
    Stack<TreeNode> predecessors = new Stack<TreeNode>();

    while(root != null){
        if(root.val >= target){
            successors.push(root);
            root = root.left;
        } else if(root.val < target) {
            predecessors.push(root);
            root = root.right;
        } 
    }

    while(k > 0 && (!successors.isEmpty() || !predecessors.isEmpty())){
        if(successors.isEmpty()){
            res.add(getNextPredecessor(predecessors));
        } else if(predecessors.isEmpty()){
            res.add(getNextSuccessor(successors));
        } else{
            if(Math.abs(predecessors.peek().val - target) >  Math.abs(successors.peek().val - target)){
                res.add(getNextSuccessor(successors));
            } else {
                res.add(getNextPredecessor(predecessors));
            }
        }
        k--;
    }
    
    return res;
}

private int getNextSuccessor(Stack<TreeNode> stack){
    TreeNode node = stack.pop();
    TreeNode next = node.right;
    while(next != null){
        stack.push(next);
        next = next.left;
    }
    return node.val;
}

private int getNextPredecessor(Stack<TreeNode> stack){
    TreeNode node = stack.pop();
    TreeNode next = node.left;
    while(next != null){
        stack.push(next);
        next = next.right;
    }
    return node.val;
}