K closest points with a size-k heap in Python

Why this matters

Many programs need the nearest things. A delivery app wants the five nearest drivers. A map app wants the ten nearest shops. A game wants the enemies closest to the player. Each case has a list of positions and one target position, and the program needs the k closest positions.

This lesson solves that problem with a heap of size k. It uses the same trick as the previous lesson: keep only the best k items seen so far. The difference is what the heap compares. It compares distances, not plain numbers. The lesson also shows the mistakes that give a wrong answer without any error message, and the ones that stop the program with an error.

Words used in this lesson

A point is a pair of whole numbers written (x, y). In Python, one point is a tuple with two items. A tuple is a fixed group of values inside parentheses, such as (3, 4). A list of points is a Python list where each item is one tuple.

The origin is the point (0, 0), the centre of the grid. Every distance in this lesson is measured from the origin.

The straight-line distance from a point to the origin is the square root of x * x + y * y. This comes from Pythagoras' rule for right triangles. The square root is slow to compute, and it gives decimal numbers.

We do not need the real distance. We only need to know which point is closer. If one squared distance is smaller than another, its square root is smaller too. So we compare the squared distance, x * x + y * y, and skip the square root. It is a whole number, which is easier to store and compare.

For example, the point (3, 4) has a real distance of 5 and a squared distance of 25. The point (1, 1) has a squared distance of 2. The point (1, 1) is closer, and both measures agree.

Python's heapq module manages a heap stored in a plain list. It only builds a min-heap, which keeps the smallest item on top. The trick for a max-heap is to store negative numbers. The next section explains how this works.

A small example first

Take the points (3, 3), (5, -1) and (-2, 4), and set k to 2. Work out the squared distance of each:

• (3, 3): 3 * 3 + 3 * 3 = 18
• (5, -1): 5 * 5 + (-1) * (-1) = 26
• (-2, 4): (-2) * (-2) + 4 * 4 = 20

The two closest points are (3, 3) at 18 and (-2, 4) at 20. The farthest point, (5, -1) at 26, is dropped. Negative coordinates work the same way. A negative number times a negative number is positive, so (-2, 4) is measured like any other point.

The idea: a max-heap of size k

We need the k closest points. Each time we hold more than k points, we should throw away the farthest one. A max-heap keeps its largest item on top. If the items are points ordered by squared distance, the top is the farthest point kept so far. That is the point to throw away.

heapq only builds a min-heap, so the code stores each point as a pair (-squared distance, point). The negative sign flips the order. The farthest point has the most negative first number, so it is the smallest pair and it sits on top. heapq.heappop removes the smallest item, which is the farthest point.

The method has three steps:

1. Add the next point to the heap with heapq.heappush. The heap keeps the pair with the most negative number on top.
2. If the heap now holds more than k points, remove the top with heapq.heappop. The top is the farthest point, so it goes.
3. Repeat for every point in the list.

At the end, the heap holds the k closest points.

Why not keep the smallest squared distance on top? A min-heap of the plain squared distances keeps the closest point on top. heappop would then remove the closest point, which is the one point we need to keep. The wrong answer comes with no error message, and the mistakes section shows it.

Here is the heap after each step for the small example. The heap is a plain list, and print shows its items in their internal order. The first item is the top.

add (3, 3) -> [(-18, (3, 3))]
add (5, -1) -> [(-26, (5, -1)), (-18, (3, 3))]
add (-2, 4) -> [(-26, (5, -1)), (-18, (3, 3)), (-20, (-2, 4))]
  pop removes (5, -1) -> [(-20, (-2, 4)), (-18, (3, 3))]

After the third point is added, the heap holds three points. That is more than k = 2, so heappop removes the top. The top is (5, -1), the farthest point with squared distance 26. The two points left are (-2, 4) and (3, 3), which are the answer.

The cost is the same as in the previous lesson. Each push or pop costs about log k steps, because the heap never holds more than k + 1 points. O(...) is the usual way to write how the work grows as the input grows. For the whole list, the cost is O(n log k) time and O(k) memory. Sorting all the points by distance would cost O(n log n) time and O(n) memory. When k is much smaller than n, the heap does less work and uses far less memory.

Python also has heapq.nsmallest(k, points, key=dist_sq), which returns the k smallest items in one call. Real code often uses it. This lesson writes the heap by hand so that each step is visible.

Example 1: the two-line problem

The first code sample has the k_closest function and two inputs. The first input has two points and k = 1. The second has three points and k = 2.

The output is:

[(-2, 2)]
[(-2, 4), (3, 3)]

The first line is correct. The squared distance of (-2, 2) is 8, and the squared distance of (1, 3) is 10. The closer point wins.

The second line holds the right two points, (3, 3) and (-2, 4). They appear in a different order from the one in the small example. The list comes out in the heap's internal order, not sorted by distance. The problem accepts the points in any order. If you need them closest first, sort the result. Exercise 2 shows how.

The line to read closely is heapq.heappush(heap, (-dist_sq(p), p)). Python compares tuples item by item. It compares the first items, which are the negative distances. It compares the points only when the first items are equal. The heap therefore keeps the largest squared distance on top.

The last line of the function, return [p for _, p in heap], is a list comprehension. It walks through the pairs, drops the distance with the name _, and keeps the point. The name _ is a common way to mark a value that is not used.

Example 2: checking against sorting on 10,000 points

The second example is bigger. It creates 10,000 random points, finds the 50 closest with the heap, and checks the answer against a simpler method. The simpler method sorts all 10,000 points by distance and takes the first 50. Sorting is slower, but the code is easy to trust.

The random numbers come from random.Random(2024). The number 2024 is a seed, a starting value for the random generator. The same seed gives the same points on every run, so the output does not change.

The output is:

points kept: 50
distances match sorting: True
farthest kept squared distance: 651688
closest squared distance: 5525

The check compares lists of distances, not lists of points. Two different points can sit at the same distance. If the heap and the sort pick different points from that tie, both answers are right. Comparing distances catches real errors and ignores harmless differences in which tied point was picked.

The code sorts the 50 heap distances before it compares them. The heap does not keep its items in order, so both lists must be sorted before they can be compared position by position.

Each heap step costs about log2(50), which is about 5.6. Sorting all the points costs about log2(10000), which is about 13.3, for each point. The heap does roughly half the work per point here. With a bigger list and a smaller k, the gap gets much bigger.

Example 3: ties at the boundary

Some inputs have two points at the same distance, and the cut-off falls between them. For example, with (1, 0) and (0, 1) and k = 1, both points are at squared distance 1. The heap keeps one of them.

The output is:

[(1, 0)]
[(1, 0)]

Both inputs keep (1, 0), even though the points were given in a different order. When two distances are equal, the heap compares the points themselves. The pop removes the smaller pair first, so the point with the larger coordinates stays.

Both answers are correct for the problem, because the problem asks for k closest points and does not say which point to pick among equal distances. Do not write a test that expects one specific point in that case. Compare distances, as the check in example 2 does.

Lesson 06 of this module showed quickselect. The same idea works here on squared distances, with average O(n) time. The code is longer, though, and the heap is the answer most people expect for this problem.

Common mistakes

Using a min-heap

The next code is wrong. It stores the plain squared distance, so the top of the heap is the closest point.

The output is:

[(-2, 4), (5, -1)]

The correct answer is [(3, 3), (-2, 4)]. The program prints without an error. It removed (3, 3), the closest point, and kept the farthest one. The fix is to store the negative distance, as the main code does. Test with a small case where you know the answer.

Using floats for the distance

A float is Python's type for decimal numbers, such as 2.5. A float is stored in a fixed number of bits, so it keeps about 15 to 16 decimal digits exactly. Beyond that, it rounds. Python integers do not round, because they grow to hold any size. Keep the squared distance as an integer.

The output is:

exact: 10000000200000001
through float: 10000000200000000
same: False

The last digit is lost. Coordinates this large are rare in small lessons, but they are common in real data. Two different squared distances can round to the same float, and then the heap can keep the wrong point. The division operator / also gives a float in Python 3, so do not divide the squares.

Forgetting the pop

The heap must shrink back to size k. If the heappop step is missing, the heap keeps every point, and the answer has the wrong length.

The output is:

returned 5 points, expected 2
[(9, 9), (4, 4), (2, 2), (1, 1), (3, 3)]

The fix is the same if len(heap) > k: heapq.heappop(heap) check inside the loop. It must run after every push. If you check the length only once at the end, the heap has already grown too big.

Items that cannot be compared

Python compares the items in a heap with <. The tuple (-d, p) compares its second item only when the first items are equal. A plain class has no < by default, so two objects with equal distances cannot be compared.

import heapq
 
 
class Point:
    def __init__(self, x, y):
        self.x = x
        self.y = y
 
 
def dist_sq(p):
    return p.x * p.x + p.y * p.y
 
 
heap = []
heapq.heappush(heap, (-dist_sq(Point(1, 0)), Point(1, 0)))
heapq.heappush(heap, (-dist_sq(Point(0, 1)), Point(0, 1)))

The program stops with this error. The lines above it show the file name and the line number. The last line names the error:

TypeError: '<' not supported between instances of 'Point' and 'Point'

Both distances are 1, so the first items tie, and the heap tries to compare the two Point objects. Python does not know which one is smaller. The fix is to store the coordinate tuple, as the main code does. The other fix is to give the class a method named __lt__, which Python calls for the < operator.

Practice

Exercise 1

The points are (1, 1), (2, 2) and (3, 3), and k is 2. Find the two closest points by their squared distances.

Show the answer

The squared distances are 1 * 1 + 1 * 1 = 2 for (1, 1), 2 * 2 + 2 * 2 = 8 for (2, 2), and 3 * 3 + 3 * 3 = 18 for (3, 3). The two smallest are 2 and 8. The answer is (1, 1) and (2, 2). The point (3, 3) is dropped.

Exercise 2

Change the function so the points come back closest first. Run it on the two inputs from example 1.

Show the answer

Sort the result with the same distance key that the heap uses. The key is the function that tells sorted which number to compare. The new function calls k_closest from example 1.

def k_closest_sorted(points, k):
    return sorted(k_closest(points, k), key=dist_sq)

Running it on [(3, 3), (5, -1), (-2, 4)] with k = 2 prints [(3, 3), (-2, 4)]. Running it on [(1, 1), (2, 2), (3, 3)] with k = 2 prints [(1, 1), (2, 2)]. Both answers are closest first.

Exercise 3

Java needs a cast to long before it squares a coordinate, to avoid overflow. Does the Python dist_sq need the same care? Why or why not?

Show the answer

No. Python integers have no fixed size, so 50000 * 50000 gives exactly 2500000000 and no cast is needed. The value 2 ** 100 (2 multiplied by itself 100 times) is 1267650600228229401496703205376, which has 31 digits, and Python prints all of them. The float version is the one that loses digits, as the mistake section shows.

Exercise 4

The list has one million points, and k is 10. How many points can be in the heap at the same time?

Show the answer

At most 11. The heap grows to k + 1 points, which is 11, and then heappop removes one. So the heap never holds more than 11 points, no matter how long the list is. This is why the memory cost is O(k) and not O(n).

Exercise 5

The wrong min-heap version runs on (3, 3), (5, -1) and (-2, 4) with k = 2. Which point does it remove first, and why?

Show the answer

It removes (3, 3). The min-heap keeps the smallest squared distance on top. The three squared distances are 18, 26 and 20, and 18 is the smallest, so the pop removes the point with 18, which is (3, 3). That leaves (5, -1) and (-2, 4), which is the wrong answer.

Summary

• A max-heap of size k keeps the k closest points, because its top is the farthest point kept so far.
• heapq builds only min-heaps, so store (-squared distance, point) to get a max-heap.
• Compare squared distances, x * x + y * y, which gives the same order as real distances without a square root.
• The method costs O(n log k) time and O(k) memory, which is less than sorting all the points when k is small.
• Keep distances as integers, because floats can round large values and change the answer.
• Check results by comparing distances, since two points at the same distance can both be correct.