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.

Example 1:
Input:
{1}
0.000000
1
Output:
[1]

我的想法

跟二分法的658. Find K Closest Elements很像，但是对于bst来说无法直接访问前一个值和后一个值，不知道改怎么解决

解答

jiuzhang solution:
用两个stack来分别存储lower结点和upper结点，以及定义move函数来寻找下一个较大值和较小值
注意move函数的写法

public class Solution {
    public List<Integer> closestKValues(TreeNode root, double target, int k) {
        List<Integer> res = new ArrayList<>();
        Stack<TreeNode> lowerStack = getStack(root, target);
        Stack<TreeNode> upperStack = getStack(root, target);
        if(lowerStack.peek().val > target) {
            moveLower(lowerStack);
        } else {
            moveUpper(upperStack);
        }
        while(k > 0) {
            //only use moveUpper() or moveLower() to operate stack
            if(lowerStack.isEmpty() && upperStack.isEmpty()) {
                return res;
            } else if(lowerStack.isEmpty()) {
                res.add(upperStack.peek().val);
                moveUpper(upperStack);
            } else if(upperStack.isEmpty()) {
                res.add(lowerStack.peek().val);
                moveLower(lowerStack);
            } else if(target - lowerStack.peek().val < upperStack.peek().val - target) {
                res.add(lowerStack.peek().val);
                moveLower(lowerStack);
            } else {
                res.add(upperStack.peek().val);
                moveUpper(upperStack);
            }
            k--;
        }
        return res;
    }
    private Stack<TreeNode> getStack(TreeNode root, double target) {
        Stack<TreeNode> stack = new Stack<>();
        while(root != null) {
            stack.push(root);
            if(root.val > target) {
                root = root.left;
            } else {
                root = root.right;
            }
        }
        return stack;
    }
    //find the upper node of current node
    private void moveUpper(Stack<TreeNode> stack) {
        TreeNode node = stack.peek();
        if(node.right == null) {
            node = stack.pop();
            while(!stack.isEmpty() && stack.peek().right == node) {
                node = stack.pop();
            }
            return;
        }
        node = node.right;
        while(node != null) {
            stack.push(node);
            node = node.left;
        }
        return;
    }
    //find the lower node of current node
    private void moveLower(Stack<TreeNode> stack) {
        TreeNode node = stack.peek();
        if(node.left == null) {
            node = stack.pop();
            while(!stack.isEmpty() && stack.peek().left == node) {
                node = stack.pop();
            }
            return;
        }
        node = node.left;
        while(node != null) {
            stack.push(node);
            node = node.right;
        }
        return;
    }
}