K closest points with a size-k heap

Why this matters

Many programs need the nearest things to a point. A delivery app shows the five restaurants closest to you. A game finds the closest enemy to a player. A sensor network picks the three nearest sensors to a fault. In each case the input is a large list of positions, and the program needs only a few of them, the closest ones.

The previous lesson showed the same shape of problem with numbers: keep the k largest values with a heap of size k. This lesson uses that idea with points. Each point is a pair of coordinates, such as (3, 4), and the "size" of a point is how far it is from the origin, which is the point (0, 0).

You will learn how to measure distance without a square root, how to keep the k closest points in a heap, and how to check the result against a simple sort.

Points and distance

A coordinate is a number that gives a position along one direction. A point on a flat map needs two coordinates, x for left and right and y for up and down. The point (3, 4) is 3 steps right and 4 steps up from the origin (0, 0).

The straight-line distance from the origin to (3, 4) is 5. It comes from the Pythagorean rule: the square root of x times x plus y times y. Here that is the square root of 3 * 3 + 4 * 4, which is the square root of 25, which is 5.

Compare squares instead of distances

Taking a square root takes extra work, and it often gives a decimal that is rounded. But we do not need the actual distance to rank points. We only need to know which point is closer. If the square of one distance is smaller than the square of another, the first distance is smaller too, as long as both distances are zero or more.

Here is an example with three points. Their distances are 5, about 1.41, and about 3.61. Their squared distances are 25, 2, and 13. The order is the same either way: the point with squared distance 2 is closest, then 13, then 25. So we store and compare the squared distance, which is always a whole number when the coordinates are whole numbers.

A struct for a point

A struct is a way to group related values under one name. Here the struct holds the two coordinates.

struct Point {
    long long x;
    long long y;
};

The type long long is a whole-number type that can hold very large values. A squared distance can be large, so long long leaves room. The reason is explained in the mistakes section.

The function below returns the squared distance of one point.

long long squaredDistance(const Point& p) {
    return p.x * p.x + p.y * p.y;
}

The const Point& means the function reads the point without copying it and without changing it.

Method one: sort every point

The simplest correct method sorts all the points by squared distance, closest first, and takes the first k.

std::sort(points.begin(), points.end(), [](const Point& a, const Point& b) {
    return squaredDistance(a) < squaredDistance(b);
});
points.resize(std::min(k, points.size()));

The lambda, which is a small function written inside the call, tells std::sort how to compare two points. The line points.resize cuts the list down to the first k points.

This takes O(n log n) steps, because sorting all n points is the expensive part. The program later uses this method as a reference, to check the faster method against it.

Method two: a max-heap of size k

The faster method keeps only the k closest points seen so far. It stores them in a max-heap, which is a heap where the largest value is always on top. Here the "largest" value is the farthest point kept.

The rule is simple. Read the next point and compute its squared distance d.

• If the heap holds fewer than k points, push the new point.
• If the heap is full and d is smaller than the top, the top is the farthest point kept. Remove it with a pop and push the new point.
• Otherwise ignore the new point. It is no closer than the k points already kept.

After every point has been read, the heap holds the k closest points. The top is the farthest of those k points, so one look at the top answers "is the new point closer than the farthest one I keep?"

A walk through example

Take the points (3, 3), (5, -1) and (-2, 4), and let k = 2. Their squared distances are 18, 26 and 20.

At the end the heap holds the squared distances 18 and 20, which belong to (3, 3) and (-2, 4). The point (5, -1) was removed because it was the farthest.

A longer example

The first walk through had only three points. Here is one with six points and k = 3. The rule is the same, but now the heap is full for most of the trip, so the "ignore" case happens too.

The heap ends with squared distances 2, 5 and 8. Those belong to (-1, 1), (1, 2) and (2, 2). Sorting the six points by squared distance gives 2, 5, 8, 9, 16, 16, and the first three match.

Notice that the point (-3, 0) was kept for a while and then removed. It was not the closest or the farthest of the points read so far, but it was the farthest of the three kept when (-1, 1) arrived. A point can enter the heap and leave it later. The heap only ever holds the best k points seen so far, so a point that leaves can never come back.

A tie at the boundary

Ties are the case where the heap needs a rule. Take three points, (2, 0), (0, 2) and (1, 1), with k = 2. The first two points are both at squared distance 4.

• (2, 0) has room, so it is pushed. The heap holds one pair, distance 4 with index 0.
• (0, 2) has room, so it is pushed. The heap now holds two pairs with distance 4. The pair with the larger index, index 1, is the bigger pair in the comparison, so it sits on top.
• (1, 1) has squared distance 2, which is less than the top distance 4. The program pops the top, which is (0, 2), and pushes (1, 1).

The heap now holds (2, 0) and (1, 1), and the program returns [[1, 1], [2, 0]]. Here the point read later, (0, 2), was removed because it sat on top. Two kept points that tie are ordered by index, so the later index sits on top and leaves first. A new point that ties the top does not replace it, because the replacement needs a strictly smaller distance. So an equal distance never pushes out a point that is already kept.

The program

Here is a complete program. The heap stores pairs of a squared distance and the index of the point. The index tells the program which point the pair belongs to.

#include <algorithm>
#include <cstddef>
#include <cstdint>
#include <iostream>
#include <queue>
#include <utility>
#include <vector>
 
struct Point {
    long long x;
    long long y;
};
 
// The squared distance from the origin. Taking the square root is not needed,
// because the square root never changes which of two values is larger.
long long squaredDistance(const Point& p) {
    return p.x * p.x + p.y * p.y;
}
 
// Returns the k points closest to the origin, closest first.
std::vector<Point> kClosest(const std::vector<Point>& points, std::size_t k) {
    if (k == 0) {
        return {};
    }
 
    // Max-heap of (squared distance, index). The top is the farthest point kept so far.
    std::priority_queue<std::pair<long long, std::size_t>> heap;
    for (std::size_t i = 0; i < points.size(); ++i) {
        const long long d = squaredDistance(points[i]);
        if (heap.size() < k) {
            heap.push({d, i});
        } else if (d < heap.top().first) {
            heap.pop();  // remove the farthest kept point
            heap.push({d, i});
        }
    }
 
    // Move the kept pairs out of the heap, then sort them so the closest comes first.
    std::vector<std::pair<long long, std::size_t>> kept;
    while (!heap.empty()) {
        kept.push_back(heap.top());
        heap.pop();
    }
    std::sort(kept.begin(), kept.end());
 
    std::vector<Point> result;
    for (const auto& [distance, index] : kept) {
        result.push_back(points[index]);
    }
    return result;
}
 
// The reference answer: sort every point, then take the first k.
std::vector<Point> kClosestBySort(std::vector<Point> points, std::size_t k) {
    std::sort(points.begin(), points.end(), [](const Point& a, const Point& b) {
        return squaredDistance(a) < squaredDistance(b);
    });
    points.resize(std::min(k, points.size()));
    return points;
}
 
void printPoints(const std::vector<Point>& points) {
    std::cout << "[";
    for (std::size_t i = 0; i < points.size(); ++i) {
        std::cout << (i ? ", " : "") << "[" << points[i].x << ", " << points[i].y << "]";
    }
    std::cout << "]\n";
}
 
int main() {
    std::vector<Point> first{{1, 3}, {-2, 2}};
    std::cout << "example 1, k=1: ";
    printPoints(kClosest(first, 1));
 
    std::vector<Point> second{{3, 3}, {5, -1}, {-2, 4}};
    std::cout << "example 2, k=2: ";
    printPoints(kClosest(second, 2));
 
    // A larger test: 200000 pseudo-random points, checked against the sort answer.
    std::uint32_t state = 7;
    auto next = [&state]() {
        state = state * 1664525u + 1013904223u;
        return static_cast<long long>(state % 20001u) - 10000;  // -10000 to 10000
    };
    std::vector<Point> big;
    for (int i = 0; i < 200000; ++i) {
        big.push_back({next(), next()});
    }
    const std::size_t k = 5;
    auto fast = kClosest(big, k);
    auto slow = kClosestBySort(big, k);
    bool same = true;
    for (std::size_t i = 0; i < k; ++i) {
        if (squaredDistance(fast[i]) != squaredDistance(slow[i])) {
            same = false;
        }
    }
    std::cout << "200000 points, k=5, distances match the sort answer: " << (same ? "yes" : "no") << '\n';
    std::cout << "closest of the large test: ";
    printPoints(fast);
    return 0;
}

Compile and run it with:

g++ -std=c++20 -Wall -o k_closest k_closest.cpp && ./k_closest

The output is:

example 1, k=1: [[-2, 2]]
example 2, k=2: [[3, 3], [-2, 4]]
200000 points, k=5, distances match the sort answer: yes
closest of the large test: [[-5, -15], [30, 25], [-29, -30], [26, 41], [-61, -13]]

The first example has two points. (1, 3) has squared distance 10, and (-2, 2) has squared distance 8. So (-2, 2) is the closest and it is the answer.

The second line matches the walk through above. The large test compares the heap answer with the sort answer on 200,000 points. The line says yes, which means the five smallest squared distances are the same in both methods. The closest points printed at the end are sorted from closest to farthest, and their squared distances rise from 250 to 3890.

Why the cost is lower

Each point costs at most one push and one pop on a heap that holds at most k points. A push or pop on a heap of size k costs O(log k) steps. So the whole method costs O(n log k), and it keeps about k points in memory.

Compare the two numbers. The base-2 logarithm of 200,000 is about 17.6, and the base-2 logarithm of 5 is about 2.3. Sorting pays about 17.6 steps for each point. The heap method pays about 2.3 steps for each point. When k is much smaller than n, that gap is the reason to use the heap.

The heap method is not always the best. When k is close to n, the heap is nearly as large as the list, and sorting is just as good. A quickselect, which was covered in the previous lesson, can also find the kth smallest squared distance first and then keep the points at or below that value. It needs the whole list in memory, while the heap method needs only k points at a time.

How ties are handled

Two points can be exactly the same distance from the origin. The heap stores pairs of (squared distance, index), so a tie is broken by the index. A new point replaces the top only when its distance is strictly smaller, so the point read first is kept when two points tie.

Here is a check with two points at the same distance and k = 1.

(1, 0) and (0, 1) both have squared distance 1.
The first point is read first, so it is kept.

The returned answer is (1, 0). The final output is sorted by distance and then by index, so the order of equal distances is the order the points were given.

Common mistakes

Using a min-heap

The heap must be a max-heap. A min-heap keeps the top as the smallest value, so when it is full the pop removes the closest point kept, not the farthest. The check then lets in only points closer than that one, so the wrong points are kept.

This is the wrong version of the function, with the same logic and a min-heap:

std::priority_queue<std::pair<long long, std::size_t>,
                    std::vector<std::pair<long long, std::size_t>>,
                    std::greater<>> heap;

Run on the second example with k = 2, the min-heap returns (3, 3) and (5, -1), with squared distances 18 and 26. It skips (-2, 4), which is closer than (5, -1). The program does not print an error. It returns the wrong points. Check the heap direction by hand on a three-point example.

Using int for squared distances

An int on most current systems holds numbers up to about 2 billion. A squared coordinate can pass that limit quickly. The following program squares the number 50000 in two ways.

#include <iostream>
 
int main() {
    int x = 50000;
    int bad = x * x;
    long long good = static_cast<long long>(x) * x;
    std::cout << "int: " << bad << '\n';
    std::cout << "long long: " << good << '\n';
    return 0;
}

Compiled with g++ -std=c++20 -Wall -Wextra -o overflow_int overflow_int.cpp, it prints:

int: -1794967296
long long: 2500000000

The correct value is 2,500,000,000. The int version printed a negative number. In C++, signed integer overflow is undefined behaviour, which means the program is not allowed to do anything specific, and the result can change with another compiler or setting. The compiler did not warn here, so the fix is a rule you apply yourself: use long long for squared distances.

Reading the heap in the wrong order

Popping a max-heap gives the largest value first. If you copy the heap into a list by popping, the list comes out farthest first. The program in this lesson reverses that by sorting the kept pairs before it builds the answer. If you skip the sort, the answer comes out in the wrong order, and a test that checks the order will fail.

Calling top on an empty heap

When k is 0, the heap never fills, so the else if line reads heap.top() on an empty heap. Reading the top of an empty heap is undefined behaviour. The early return if (k == 0) return {}; at the start of the function prevents this.

Using a square root for no reason

Writing std::sqrt on each squared distance does not give a wrong answer. It only adds work for each point, and the decimal values can be rounded. Compare the squared values directly.

Practice

Exercise 1. Use the heap method with k = 2 on the points (1, 1), (2, 2), (3, 0) and (0, 5). Which points are returned, and in what order?

Show the answer

The squared distances are 2, 8, 9 and 25. The heap takes (1, 1) and (2, 2) first, since it has room. The point (3, 0) has 9, which is not less than the top 8, so it is ignored. The point (0, 5) has 25, which is also ignored. The answer is (1, 1) and then (2, 2), closest first.

Exercise 2. Change the program so that it returns the k farthest points instead of the k closest. Which two lines change?

Show the answer

Use a min-heap, so the top is the smallest value kept, and remove the closest point when the heap is full. The type of the heap becomes a std::priority_queue of the pair type with std::greater<> as the third argument, as shown in the mistakes section. The test on the top changes from d < heap.top().first to d > heap.top().first. Everything else stays the same.

Exercise 3. Why does the program compare squared distances and not real distances?

Show the answer

Squaring keeps the order of non-negative numbers, so the closest point is the same either way. Squared distances of whole-number coordinates are whole numbers, so the comparison is exact. A square root would add work for every point and might round the result.

Exercise 4. The points (1, 0) and (0, 1) are both at distance 1. Which one does kClosest return for k = 1, and why?

Show the answer

It returns (1, 0). That point is read first, and it is stored in the heap with index 0. The second point has the same squared distance, 1, and the test d < heap.top().first is false, so the heap keeps the first point.

Exercise 5. A list has 1,000,000 points and you need the 3 closest. Roughly how many steps per point does each method take: sorting, and the heap of size k? Give the numbers in terms of the base-2 logarithm.

Show the answer

Sorting takes about log2 of 1,000,000 steps per point, which is about 20. The heap method takes about log2 of 3 steps per point, which is about 1.6. These are rough counts for comparing the two methods, not measured times.

Summary

• Compare squared distances instead of real distances. The order is the same, and whole-number coordinates give exact whole-number results.
• Keep the k closest points in a max-heap of size k. Its top is the farthest kept point, and a new point replaces it only when the new point is closer.
• The heap method costs O(n log k) and keeps about k points in memory. Sorting costs O(n log n).
• Use a max-heap, not a min-heap. A min-heap keeps the farthest points instead.
• Store squared distances in long long, because an int overflows for large coordinates.
• Check a new version of the function against a simple sort on random data before you trust it.