Closest Binary Search Tree Value II
Problem
Given the root of a BST and a target value, return the k values that are closest to the target.
root = [4,2,5,1,3], target = 3.714, k = 2[4,3]def closest_k_values(root, target, k):
from collections import deque
out = deque()
def inorder(node):
if not node: return
inorder(node.left)
if len(out) < k:
out.append(node.val)
elif abs(node.val - target) < abs(out[0] - target):
out.popleft(); out.append(node.val)
else:
return # pruning is possible here
inorder(node.right)
inorder(root)
return list(out)function closestKValues(root, target, k) {
const out = [];
function inorder(node) {
if (!node) return;
inorder(node.left);
if (out.length < k) out.push(node.val);
else if (Math.abs(node.val - target) < Math.abs(out[0] - target)) {
out.shift(); out.push(node.val);
} else return;
inorder(node.right);
}
inorder(root);
return out;
}class Solution {
Deque out = new ArrayDeque<>();
int K; double T;
public List closestKValues(TreeNode root, double target, int k) {
K = k; T = target;
inorder(root);
return new ArrayList<>(out);
}
void inorder(TreeNode node) {
if (node == null) return;
inorder(node.left);
if (out.size() < K) out.add(node.val);
else if (Math.abs(node.val - T) < Math.abs(out.peekFirst() - T)) {
out.pollFirst(); out.add(node.val);
} else return;
inorder(node.right);
}
} class Solution {
deque out;
int K; double T;
void inorder(TreeNode* node) {
if (!node) return;
inorder(node->left);
if ((int)out.size() < K) out.push_back(node->val);
else if (abs(node->val - T) < abs(out.front() - T)) {
out.pop_front(); out.push_back(node->val);
} else return;
inorder(node->right);
}
public:
vector closestKValues(TreeNode* root, double target, int k) {
K = k; T = target;
inorder(root);
return vector(out.begin(), out.end());
}
}; Explanation
We want the k values closest to a target inside a BST. The neat trick is that an in-order traversal of a BST visits the values in sorted order, so we can scan them smallest to largest while keeping only the best k seen so far.
We keep a sliding window (a deque called out) of at most k values. While the window is not yet full, we just add each value. Once it has k values, we compare the new value with the oldest/front value in the window.
Because values arrive in increasing order, the front of the window is always the value farthest below the target. If the new value is closer to the target than that front value (abs(node.val - target) < abs(out[0] - target)), we drop the front with popleft() and push the new one. If it is not closer, every later value will be even farther, so we can stop.
Example: BST gives sorted order [1, 2, 3, 4, 5], target = 3.714, k = 2. Fill window with 1, 2. At 3, it beats front 1, so window becomes [2, 3]. At 4, it beats front 2, so window becomes [3, 4]. At 5 it is not closer than 3, so we stop with answer [3, 4].
This works in one pass because the sorted order lets us throw away the worst candidate the moment a better one shows up.