Closest Binary Search Tree Value II

Problem: Closest Binary Search Tree Value II

We can use a stack to achieve in-order traversal and use a heap to store possible k values and use difference between value and target as key. If we put the value with the biggest difference the top of the heap. Once there's a new closer one, we can pop that one from the heap and insert the new one.

Code in Python:

class Solution(object):
    def closestKValues(self, root, target, k):
        """
        :type root: TreeNode
        :type target: float
        :type k: int
        :rtype: List[int]
        """
        if not root: return []
        heap, stack = [], [root]
        while stack:
            if stack[-1].left:
                stack.append(stack[-1].left)
                stack[-2].left = None
            else:
                cur = stack.pop()
                if cur.right: stack.append(cur.right)
                if len(heap) < k:
                    heapq.heappush(heap, (-abs(cur.val-target), cur.val))
                elif -abs(cur.val-target) > heap[0][0]:
                    heapq.heappop(heap)
                    heapq.heappush(heap, (-abs(cur.val-target), cur.val))
        return [a[1] for a in heap]