Python solution, O(k + (N - k)*log(k)) time


  • 0
    class Solution(object):
        def closestKValues(self, root, target, k):
            def visit(root, target, count, heap):
                if not root:
                    return
                
                visit(root.left, target, count, heap)
                
                count[0] += 1
                diff = abs(target - root.val)
                if count[0] <= k:
                    heap.append((-diff, root.val))
                    if count[0] == k:
                        heapq.heapify(heap)
                else:
                    if heap and diff < -heap[0][0]:
                        heapq.heappop(heap)
                        heapq.heappush(heap, (-diff, root.val))
                
                visit(root.right, target, count, heap)
    
            heap = []
            visit(root, target, [0], heap)
            return [heapq.heappop(heap)[1] for _ in range(k)]

  • 0

    No need for a heap and its extra log(k).


  • 0

    I know, this is a preliminary solution just to start off.


Log in to reply
 

Looks like your connection to LeetCode Discuss was lost, please wait while we try to reconnect.