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

    public List closestKValues(TreeNode root, double target, int k) {
        List res = new ArrayList<>();
        Stack s1 = new Stack<>(); 
        Stack s2 = new Stack<>();
        get_predecessor(root, target, s1);
        get_successor(root, target, s2);

        while (k-- > 0) {
            if (s1.isEmpty()) res.add(s2.pop());
            else if (s2.isEmpty()) res.add(s1.pop());
            else if (Math.abs(s1.peek() - target) < Math.abs(s2.peek() - target)) res.add(s1.pop());
            else res.add(s2.pop());
        }
        return res;
    }
    
    private void get_predecessor(TreeNode root, double target, Stack s1) {
        if (root == null) return;
        get_predecessor(root.left, target, s1);
        if (root.val > target) return;
        s1.push(root.val);
        get_predecessor(root.right, target, s1);
    }
    
    private void get_successor(TreeNode root, double target, Stack s2) {
        if (root == null) return;
        get_successor(root.right, target, s2);
        if (root.val <= target) return;
        s2.push(root.val);
        get_successor(root.left, target, s2);
    }