Closest Binary Search Tree Value II

Problem

Given the root of a binary search tree, a target value, and an integer k, return the k values in the BST that are closest to the target. You may return the answer in any order.

You are guaranteed to have only one unique set of k values in the BST that are closest to the target.

Examples

Example 1

graph TD;
D(4) --- B(2) & E(5)
B --- A(1) & C(3)

Input: root = [4,2,5,1,3], target = 3.714286, k = 2
Output: [4, 3]

Follow up

Assume that the BST is balanced, could you solve it in less than O(n) runtime (where n = total nodes)?

Solution

Method 1 - Inorder Traversal + Maxheap

This problem can be efficiently solved using an in-order traversal of the Binary Search Tree (BST), which gives us elements in sorted order. By leveraging a max-heap to always keep track of the k closest elements, we can ensure a valid solution.

Approach

1. Perform an in-order traversal of the BST to explore nodes in sorted order.
2. Use a max-heap (or simply a deque in Python) to store the closest k elements based on their absolute difference from the target.
3. For each node:
   • Compute the absolute difference between the node's value and the target.
   • If the heap contains fewer than k elements, add the current node's value to the heap.
   • If the heap is full (size = k) and the current node's difference is smaller than the maximum difference in the heap, replace the largest element in the heap with the current node's value.
4. Return the contents of the heap after traversal.

Code

Java

class Solution {

    public List<Integer> closestKValues(TreeNode root, double target, int k) {
        PriorityQueue<Integer> heap = new PriorityQueue<>((a, b) ->
            Double.compare(Math.abs(b - target), Math.abs(a - target))
        );
        helper(root, target, k, heap);
        return new ArrayList<>(heap);
    }

    private void helper(
        TreeNode node,
        double target,
        int k,
        PriorityQueue<Integer> heap
    ) {
        if (node == null) return;

        helper(node.left, target, k, heap);

        // Insert the current value into the heap
        if (heap.size() < k) {
            heap.add(node.val);
        } else if (
            Math.abs(node.val - target) < Math.abs(heap.peek() - target)
        ) {
            heap.poll();
            heap.add(node.val);
        }

        helper(node.right, target, k, heap);
    }
}

Python

class Solution:
    def closestKValues(self, root: Optional[TreeNode], target: float, k: int) -> List[int]:
        ans: deque[int] = deque()
        
        def dfs(node: Optional[TreeNode]) -> None:
            if not node:
                return
            
            dfs(node.left)
            
            # Use deque to maintain k closest values
            if len(ans) < k:
                ans.append(node.val)
            elif abs(node.val - target) < abs(ans[0] - target):
                ans.popleft()
                ans.append(node.val)
            
            dfs(node.right)
        
        dfs(root)
        return list(ans)

Complexity

• ⏰ Time complexity: O(n log k)

  • In-order traversal: O(n) for traversing all nodes.
  • Heap operations for k nodes: O(log k) per insertion, which happens for every node. Thus, total heap operations = O(n log k).

• 🧺 Space complexity: O(h + k)

  • In-order traversal stack + heap:
    • Stack requires O(h) space for recursion (where h is the height of the BST, O(log n) for balanced BST and O(n) for skewed BST).
    • Max heap stores k elements, so O(k).