Python with heapq in O(N lgk)


  • 0
    A
    def closestKValues(self, root, target, k):
    
        heap = [] # min heap
        stack = []
        while root or stack:
            if root:
                stack.append(root)
                root = root.left
            else:
                tmpNode = stack.pop()
                root = tmpNode.right
                if len(heap) < k:
                    heapq.heappush(heap, (- abs(target - tmpNode.val), tmpNode.val) )
                elif -abs(target - tmpNode.val) > heap[0][0]:
                    heapq.heappop(heap)
                    heapq.heappush(heap, (-abs(target - tmpNode.val), tmpNode.val))
        
        return [a[1] for a in heap]

Log in to reply
 

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