Closest Binary Search Tree Value II

这道题是在二叉树上做 range query.

二叉树上做 range query 普遍要依靠 getPrev() 和 getNext() 函数，或者利用 Parent 指针做 traversal.

但是 BST 不一样，如上题所说， inorder 是 BST 的灵魂。因此这道题可以分为三部分；+

- 做出 inorder traversal list
- 找 target 对应的 index
- 从 index 两边扫，寻找 k 个元素

只在主函数的起始 index 上出了一个小 bug，基本算是一次 AC.

时间复杂度 O(n) + O(log n) + O(k) = O(n)

public class Solution {
    public List<Integer> closestKValues(TreeNode root, double target, int k) {

        List<Integer> inorder = inorder(root);
        int index = binarySearch(inorder, target);
        List<Integer> list = new ArrayList<Integer>();

        int start = index - 1;
        int end = index;
        while(start >= 0 && end < inorder.size()){
            int num = (Math.abs(inorder.get(start) - target) < Math.abs(inorder.get(end) - target)) 
                      ? inorder.get(start--)
                      : inorder.get(end++);
            list.add(num);
            if(list.size() == k) return list;
        }
        while(start >= 0){
            list.add(inorder.get(start--));
            if(list.size() == k) return list;
        }
        while(end < inorder.size()){
            list.add(inorder.get(end++));
            if(list.size() == k) return list;
        }

        return list;
    }

    private List<Integer> inorder(TreeNode root){
        List<Integer> list = new ArrayList<Integer>();
        TreeNode cur = root;
        while(cur != null){
            if(cur.left == null){
                list.add(cur.val);
                cur = cur.right;
            } else {
                TreeNode prev = cur.left;
                while(prev.right != null && prev.right != cur){
                    prev = prev.right;
                }
                if(prev.right == null){
                    prev.right = cur;
                    cur = cur.left;
                } else {
                    prev.right = null;
                    list.add(cur.val);
                    cur = cur.right;
                }
            }
        }
        return list;
    }

    private int binarySearch(List<Integer> list, double target){
        int start = 0;
        int end = list.size() - 1;
        while(start + 1 < end){
            int mid = start + (end - start) / 2;
            if(list.get(mid) == target){
                return mid;
            } else if(target < list.get(mid)){
                end = mid;
            } else {
                start = mid;
            }
        }

        if(target < list.get(start)) return start;
        if(target > list.get(end)) return end;

        return (Math.abs(list.get(start) - target) < Math.abs(list.get(end) - target)) ? start: end;
    }
}

还有一种写法, getPre和getSucc是核心, 原理是用两个stack来存pre和succ

• 先扫一遍将pre和succ压入stack

• 当pre和succ不为空时, 比较pre.peek()-target和succ.peek()-target, 选定从pre和succ弹出peek元素后, 如果是pre, 将pre.左子树的最右分支压入pre; 如果是succ, 将succ.右子树的最左分支压入succ

• 当pre和succ一方为空时, 仅从另一方找元素.

• pre和succ都为空时或者k==0时截止

时间复杂度是O(n), 当是skewed left/right tree时

注意一个问题, target的value不能和树中节点相等,不然无法判断向左边压还是右边压

public class Solution {
    public List<Integer> closestKValues(TreeNode root, double target, int k) {
        Stack<TreeNode> succ = new Stack<TreeNode>();
        Stack<TreeNode> pre = new Stack<TreeNode>();
        TreeNode mover = root;
        while(mover!=null){
            if(mover.val>=target){
                succ.push(mover);
                mover = mover.left;
            }
            else{
                pre.push(mover);
                mover = mover.right;
            }
        }
        List<Integer> ans = new ArrayList<Integer>();
        while(k>0 && (!succ.isEmpty() || !pre.isEmpty())){
            if(succ.isEmpty())
                ans.add(getNextPre(pre));
            else if(pre.isEmpty()){
                ans.add(getNextSucc(succ));
                System.out.println(succ.size());
            }
            else if(Math.abs(pre.peek().val-target)<Math.abs(succ.peek().val-target))
                ans.add(getNextPre(pre));
            else
                ans.add(getNextSucc(succ));

            k--;
        }
        System.out.println(succ.size());
        System.out.println(pre.size());
        return ans;
    }

    public int getNextPre(Stack<TreeNode> pre){
        TreeNode top = pre.pop();
        TreeNode p = top.left;
        while(p!=null){
            pre.push(p);
            p=p.right;
        }
        return top.val;
    }

    public int getNextSucc(Stack<TreeNode> succ){
        TreeNode top = succ.pop();
        TreeNode p = top.right;
        while(p!=null){
            succ.push(p);
            p = p.left;
        }
        System.out.println(top.val);
        return top.val;
    }

}

LeetCode 论坛上还有一些其他的写法，也很有意思。这个写法是 O(n log n) + O(k log n)，用自定义的 min heap.

public List<Integer> closestKValues(TreeNode root, final double t, int k) {
    List<Integer>ret = new ArrayList<Integer>();
    PriorityQueue<TreeNode> queue = new PriorityQueue<TreeNode>(new Comparator<TreeNode>(){
       public int compare(TreeNode n1, TreeNode n2) {
           return Math.abs(n1.val - t) < Math.abs(n2.val - t) ? -1 : 1;
       }
    });
    findClosest(root, queue);
    while(k-- > 0) ret.add(queue.poll().val);
    return ret;
}

public void findClosest(TreeNode root, PriorityQueue<TreeNode> queue) {
    if(root == null) return;
    findClosest(root.left, queue);
    queue.add(root);
    findClosest(root.right, queue);
}

这个解法我很喜欢，类似于 sliding window maximum 一样，维护一个大小为 k 的 sliding window，在树上做 inorder traversal (单调递增)，每当 window size == k 但是新 node 值比 head 的值更接近于 target 的时候，就给替换掉。当后来发现新来的 node 不比 head 更接近于 target 的时候，就可以返回结果了，而且因为 LinkedList 继承了 List，连类型转换都不需要。

时间复杂度 O(n) ，还带 early termination，非常简洁高效的解法，比我写的那个速度快，而且简洁明了。

public List<Integer> closestKValues(TreeNode root, double target, int k) {
    List<Integer> res = new LinkedList<Integer>();
    helper(root, target, k, res);
    return res;
}
private void helper(TreeNode root, double target, int k, List<Integer> res) {
    if (root == null) {
        return;
    }
    helper(root.left,target,k,res);
    if (res.size()< k) {
        res.add(root.val);
    } else {
        if (Math.abs(res.get(0)-target) > Math.abs(root.val-target)) {
            res.remove(0);
            res.add(root.val);
        } else {
            return; //early termination
        }
    }
    helper(root.right,target,k,res);
}