272 Closest Binary Search Tree Value II

Given a non-empty binary search tree and a target value, findkvalues in the BST that are closest to the target.

Note:

• Given target value is a floating point.

• You may assume k is always valid, that is: k ≤ total nodes.

• 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)?

The Idea: Use a priority queue to maintain the top k closest values (similarly to the k closest distances problem). There are a few tricks I did to get this to work with a priority queue. Since by default the priority queue will compare the first element of a tuple, our main comparator will be defined as distance(target, root.val). Secondly because priority queue defaults as a min heap, and we need to treat it as a max heap, and inverse the values with (-), and then reverse the logic we would have defined for a min heap. Secondly, because there may be duplicates in our heap (that are equidistant to the target), we need a second parameter in our tuple that will resolve this arbitrarily. I resolved this by passing in a unique counter id for each node.

Complexity: O(Nlog(k)) time since we iterate through all the elements in the tree, processing each element will take time proportional to the number of elements in the pq. O(K) space.

from queue import PriorityQueue


class Solution:
    def closestKValues(self, root, target, k):
        """
        :type root: TreeNode
        :type target: float
        :type k: int
        :rtype: List[int]
        """

        fix_same_dist_conflict = [0]
        pq = PriorityQueue()
        def pre_order_dfs(root):
            if root:
                # if the current distance exceeds the current
                # smallest maximum of the pq (inverse logic)
                cur_dist = -abs(target - root.val)
                if pq.qsize() == k and pq.queue[0][0] < cur_dist:
                    pq.get()
                    pq.put((cur_dist, fix_same_dist_conflict[0], root))
                elif pq.qsize() < k:
                    pq.put((cur_dist, fix_same_dist_conflict[0], root))
                fix_same_dist_conflict[0] += 1
                pre_order_dfs(root.left)
                pre_order_dfs(root.right)

        pre_order_dfs(root)
        solution = []
        while not pq.empty():
            solution.append(pq.get()[2].val)
        return solution