题目地址:
https://www.lintcode.com/problem/901/description
给定一棵二叉搜索树,再给定一个数值 x x x和一个非负整数 k k k,求树中与 x x x最近的 k k k个数的数值。题目保证 k k k合法。
当在一个数组里求这 k k k个数的时候是很容易的,只需找出离 x x x最近的那个数,然后向左向右取前 k k k个最近的数即可。现在是在一棵BST里做,可以先类似地寻找离 x x x最近的链(这里的链指的是从树根出发的一条路径),然后在链上向左右扩展。具体来说,先从树根向下走,设当前位置的数值是 c c c,那么当 c > x c>x c>x的时候向左走,否则向右走,同时开两个栈,分别存这条链上比 x x x小和比 x x x大的节点(之所以开栈,是在模仿非递归中序遍历一棵二叉树),这样走完之后,两个栈其一的栈顶一定就是离 x x x最近的数。接着就可以模仿一维数组里的方法了,向两边扩展,扩展完一个数就求一下前驱或者后继,将离 x x x最近的 k k 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 ) O(h+k) O(h+k), h h h是树高, 每次move的时候均摊复杂度是 O ( 1 ) O(1) O(1),空间 O ( h ) O(h) O(h)。