Heap · Top-k with a max-heap of size k

Medium · O(n log k) time · O(k) space

Problem

You are given an array points where points[i] = [x, y] is a point on the plane, and an integer k. Return the k points closest to the origin (0, 0), using straight-line (Euclidean) distance. The answer may be in any order, and it is guaranteed to be unique.
Example 1 Input: points = [[1, 3], [-2, 2]], k = 1 Output: [[-2, 2]] (distances √10 and √8, so (-2, 2) is closer) Example 2 Input: points = [[3, 3], [5, -1], [-2, 4]], k = 2 Output: [[3, 3], [-2, 4]] (distances √18, √26, √20: drop the farthest)
Constraints 1 ≤ k ≤ points.length ≤ 10^4 -10^4 ≤ x, y ≤ 10^4

Brute force

You have a list of points on a map, and you stand at the centre, the origin. Which k points are closest to you?

  1. The distance of a point (x, y) is the straight-line distance √(x² + y²).
  2. Keep exactly k points: the ones with the smallest distance.
  3. The order of the answer does not matter.

Two ideas help. First, we never need the square root. If one number is smaller than another, so is its square root, so we can compare x² + y² and stay with whole numbers. Second, the answer is just a ranking: a point is in the answer only if fewer than k points are closer. Here are the squared distances of five points, with k = 2.

The direct idea: sort all points by distance and take the first k.

The problem with it

Sorting takes about n log n steps. If k is 1, that is a lot of work to find one smallest point. And it puts all n points in order when we only care about k of them.

How to think your way to the optimal solution

The sorted list is a ranking, and we throw away everything after position k. So what we really need is a way to always know the worst point we are currently keeping. While scanning the points, how do we find the farthest of our k keepers in an instant?

Turn the goal around. Instead of choosing the k closest, we keep a group of at most k points and throw out the farthest whenever the group gets too big. A heap is a box that always hands us the biggest (or smallest) item first, and removes it cheaply. We want the biggest distance first, so we use a "max-heap".

Python's heapq only has a min-heap, so we store the distance as a negative number: the smallest negative number is the biggest distance.

Why throwing out the top is safe. When the heap holds k + 1 points, its top is farther than the other k. So at least k points are closer than it, and it cannot be among the k closest. Removing it loses nothing. Note that the point we throw out may be the very one we just added.

The optimal solution

  1. Start with an empty heap.
  2. For each point, compute dist = x² + y² and add (−dist, x, y) to the heap.
  3. the heap has more than k points⟶take out the top: the farthest of the k + 1
  4. When all points are done, the heap holds the k closest points.

By hand, with the points (3,3) (1,1) (5,−1) (0,2) (−2,4) and k = 2. Their squared distances are 18, 2, 26, 4 and 20:

  1. (3,3), distance 18: add it. The heap holds {18}.
  2. (1,1), distance 2: add it. The heap holds {18, 2}. It has exactly k points.
  3. (5,−1), distance 26: add it, and now there are 3. The top is the farthest, 26: take it out. The heap holds {18, 2}.
  4. (0,2), distance 4: add it. Three points again. The farthest is 18: take it out. The heap holds {2, 4}.
  5. (−2,4), distance 20: add it. The farthest is this very point, 20: take it out. The heap holds {2, 4}.

The answer is the points with distances 2 and 4: (1,1) and (0,2). Press Play to watch the heap. The red top is the point that gets taken out when the heap grows past k.

Complexity analysis

O(n log k)
Time. Here's why: there are n points. Adding to a heap that holds at most k + 1 points takes about log k steps, and so does taking out the top. So each point costs about log k, and n points cost n log k. That is smaller than the n log n of sorting when k is small.
O(k)
Space. The heap never holds more than k + 1 points. The answer is those k points.

Interviewer follow-ups

These are the questions that usually come next. Open one to see an answer you could say out loud.

"The points arrive one at a time, forever. Can you always tell me the k closest so far?"

What they are checkingWhether you see that the heap solution never needs the whole list. This is where it clearly beats sorting.

A strong answerYes. The heap is already built one point at a time: a new point is added, and if the heap is too big, the farthest is taken out. At any moment, the heap holds the k closest so far. It uses O(k) memory no matter how many points have arrived, and each new point costs O(log k).

"Can you beat O(n log k)?"

What they are checkingWhether you know the other standard tool for "the k smallest", and what it trades away.

A strong answerYes, with quickselect. It works like quicksort, but after splitting the points around a pivot distance, it continues on only the side that contains position k. On average that takes O(n) time. The costs: it needs all the points in memory at once, it changes the order of the list, and its worst case is O(n²) unless the pivot is chosen at random. So I would use the heap for a stream or a very large list that does not fit in memory, and quickselect when everything is in memory and I want the fastest average.

"What if k is almost as big as n?"

What they are checkingWhether you notice that the cost depends on k, and can change the question when k is large.

A strong answerWhen k is close to n, the heap is nearly as big as the input, so n log k is nearly n log n and the heap gives little. I can flip the question: find the n − k farthest points instead, using a small heap of size n − k, and return everything else. Or I can simply sort, or use quickselect. I would pick whichever has the smaller heap.

Spot this pattern