Closest Binary Search Tree Value II

Tested Python solution for LeetCode 272 with 14 pytest cases. Generate a practice environment with lcpy.

LeetCode 272, Hard. Topics: Stack, Tree, Depth-First Search, Binary Search Tree, Two Pointers, Binary Tree, Heap (Priority Queue). View on LeetCode.

Generate this problem as a practice environment: tested reference solution, 14 parametrized pytest cases, and a playground notebook:

lcpy gen -n 272   # by problem number
lcpy gen -s closest_bst_value_ii   # by problem name

​

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.

Follow up: Assume that the BST is balanced. Could you solve it in less than O(n) runtime (where n = total nodes)?

​

Examples

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

Input: root = [1], target = 0.000000, k = 1
Output: [1]

​

Constraints

• The number of nodes in the tree is n.
• 1 <= k <= n <= 10^4.
• 0 <= Node.val <= 10^9
• -10^9 <= target <= 10^9

​

Solution

Reference implementation from solution.py on GitHub, full suite in test_solution.py:

import heapq

from leetcode_py import TreeNode


class Solution:
    # Time: O(n log k) where n is the node count
    # Space: O(h + k)
    def closest_k_values(self, root: TreeNode[int], target: float, k: int) -> list[int]:
        heap: list[tuple[float, int]] = []
        stack: list[TreeNode[int]] = []
        node: TreeNode[int] | None = root
        while stack or node is not None:
            while node is not None:
                stack.append(node)
                node = node.left
            node = stack.pop()
            diff = abs(node.val - target)
            if len(heap) < k:
                heapq.heappush(heap, (-diff, node.val))
            elif -heap[0][0] > diff:
                heapq.heapreplace(heap, (-diff, node.val))
            node = node.right
        return [val for _, val in heap]

​

Complexity

Time	Space
O(n log k) where n is the node count	O(h + k)

​

Tags

NeetCode All.