all topics

Heap / Priority Queue

Find K Closest Elements

Find the k closest values to x in a sorted array arr, return them in sorted order. arr = [1,2,3,4,5], k = 4, x = 3 → [1,2,3,4].

Max-heap of size k

  • Push every element onto a max-heap, keyed by how far it is from x (negated, since Python's heapq is a min-heap).
  • Heap grows past size k → pop the farthest element off, so only the k closest ever survive.
  • Drain the heap into the answer, then sort it, the heap doesn't come out in order.
python
class Solution:
    def findClosestElements(self, arr: List[int], k: int, x: int) -> List[int]:
        pq = []
 
        for ele in arr:
            heapq.heappush(pq, (-abs(x - ele), -ele))
            if len(pq) > k:
                heapq.heappop(pq)
        ans = []
        while pq:
            ans.append(-heapq.heappop(pq)[1])
        ans.sort()
        return ans
timeO(n log k)spaceO(k)