【Lintcode】901. Closest Binary Search Tree Value II

题目地址：

https://www.lintcode.com/problem/901/description

给定一棵二叉搜索树，再给定一个数值$x$和一个非负整数$k$，求树中与$x$最近的$k$个数的数值。题目保证$k$合法。

当在一个数组里求这$k$个数的时候是很容易的，只需找出离$x$最近的那个数，然后向左向右取前$k$个最近的数即可。现在是在一棵BST里做，可以先类似地寻找离$x$最近的链（这里的链指的是从树根出发的一条路径），然后在链上向左右扩展。具体来说，先从树根向下走，设当前位置的数值是$c$，那么当$c>x$的时候向左走，否则向右走，同时开两个栈，分别存这条链上比$x$小和比$x$大的节点（之所以开栈，是在模仿非递归中序遍历一棵二叉树），这样走完之后，两个栈其一的栈顶一定就是离$x$最近的数。接着就可以模仿一维数组里的方法了，向两边扩展，扩展完一个数就求一下前驱或者后继，将离$x$最近的$k$个数找到。代码如下：

import java.util.ArrayList;
import java.util.Deque;
import java.util.LinkedList;
import java.util.List;

public class Solution {
    /**
     * @param root:   the given BST
     * @param target: the given target
     * @param k:      the given k
     * @return: k values in the BST that are closest to the target
     */
    public List<Integer> closestKValues(TreeNode root, double target, int k) {
        // write your code here
        List<Integer> res = new ArrayList<>();
        if (root == null || k == 0) {
            return res;
        }
        
        Deque<TreeNode> stkLeft = new LinkedList<>(), stkRight = new LinkedList<>();
        fillStk(root, stkLeft, stkRight, target);
        
        for (int i = 0; i < k; i++) {
            if (stkLeft.isEmpty()) {
                res.add(stkRight.peek().val);
                moveRight(stkRight);
            } else if (stkRight.isEmpty()) {
                res.add(stkLeft.peek().val);
                moveLeft(stkLeft);
            } else {
                int valLeft = stkLeft.peek().val, valRight = stkRight.peek().val;
                if (target - valLeft <= valRight - target) {
                    res.add(valLeft);
                    moveLeft(stkLeft);
                } else {
                    res.add(valRight);
                    moveRight(stkRight);
                }
            }
        }
        
        return res;
    }
    
    // 把stk.peek()的后继所在的左链入栈，并正好让后继位于栈顶
    private void moveRight(Deque<TreeNode> stk) {
        TreeNode cur = stk.pop();
        cur = cur.right;
        while (cur != null) {
            stk.push(cur);
            cur = cur.left;
        }
    }
    
    private void moveLeft(Deque<TreeNode> stk) {
        TreeNode cur = stk.pop();
        cur = cur.left;
        while (cur != null) {
            stk.push(cur);
            cur = cur.right;
        }
    }
    
    private void fillStk(TreeNode root, Deque<TreeNode> stkLeft, Deque<TreeNode> stkRight, double target) {
        while (root != null) {
            if (target < root.val) {
                stkRight.push(root);
                root = root.left;
            } else {
                stkLeft.push(root);
                root = root.right;
            }
        }
    }
}

class TreeNode {
    int val;
    TreeNode left, right;
    
    TreeNode(int x) {
        val = x;
    }
}

时间复杂度$O(h+k)$，$h$是树高， 每次move的时候均摊复杂度是$O(1)$，空间$O(h)$。