Start learning free

LESSON 384 OF 535 · HEAP

K Closest Points to Origin in C++: O(n log k) with a Max-Heap

Practice O(n log k) Read the full explanation

Given a vector of 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).

The distance between a point [x, y] and the origin is the usual straight-line distance, √(x2 + y2).

Return the points sorted by their distance to the origin, the closest first. Points at the same distance come in order of their x, and points with the same x too in order of their y. The answer is guaranteed to be unique: every point left out is farther from the origin than every point in the answer.

 

Example 1:

Input: points = [[1,3],[-2,2]], k = 1
Output: [[-2,2]]
Explanation: [1,3] is √10 away from the origin and [-2,2] is √8 away.
√8 is smaller, so [-2,2] is the closest point.

Example 2:

Input: points = [[3,3],[5,-1],[-2,4]], k = 2
Output: [[3,3],[-2,4]]
Explanation: The squared distances are 18, 26 and 20. The two closest points are [3,3] (18) and [-2,4] (20), in that order.

Example 3:

Input: points = [[0,2],[2,0],[1,1],[0,-2],[3,3]], k = 4
Output: [[1,1],[0,-2],[0,2],[2,0]]
Explanation: [1,1] is the closest (√2). [0,2], [2,0] and [0,-2] are all 2 away, so they come by x:
[0,-2] and [0,2] (x = 0) before [2,0], and [0,-2] before [0,2] because -2 < 2. [3,3] is the farthest and is left out.

 

Constraints:

  • 1 <= k <= points.length <= 105
  • -104 <= x, y <= 104
  • The same point can appear more than once.
  • The answer is unique.

Note:

Your algorithm should run in O(n log k) time and use O(k) extra space, where n is the number of points. Sorting all the points by distance would take O(n log n).

No square roots are needed: one point is closer than another exactly when its x2 + y2 is smaller, and that squared distance is a whole number of at most 2 · 108, which fits in a 32-bit integer.

STEP 1 OF 14
    Ctrl + Enter to check

    C++ · HEAP

    K Closest Points To Origin in C++, explained

    Problem overview

    The function kClosest receives a vector<vector<int>> named points, where each points[i] is {x, y}, and an int named k. It returns the k points closest to the origin (0, 0), as a vector<vector<int>> sorted by distance, closest first. Points at the same distance come in order of x, and points with the same x in order of y. The answer is unique: every point left out is farther away than every point in the answer.

    For points = {{1, 3}, {-2, 2}} and k = 1 the answer is [[-2,2]]: {1, 3} is √10 from the origin and {-2, 2} is √8. For {{3, 3}, {5, -1}, {-2, 4}} and k = 2 the squared distances are 18, 26 and 20, so the answer is [[3,3],[-2,4]]. With {{0, 2}, {2, 0}, {1, 1}, {0, -2}, {3, 3}} and k = 4, the answer is [[1,1],[0,-2],[0,2],[2,0]]: the three points at distance 2 are ordered by x, then y, and {3, 3} is left out.

    The statement asks for O(n log k) time and O(k) extra space, where n is the number of points. Sorting all the points by distance would take O(n log n).

    Approach: keep the k best in a max-heap

    The distance to the origin is √(x² + y²), but no square root is needed. A point is closer than another exactly when its squared distance is smaller, and the squared distance is a whole number. With coordinates up to 10,000 in absolute value, it is at most 200,000,000, which fits in an int.

    To keep only the best k points while scanning, use a max-heap of the points kept so far. std::priority_queue is a max-heap by default: its top is the largest element, here the farthest point kept. Each new point goes in, and when the heap holds more than k points the top is removed. Out of k + 1 points the farthest can't be among the k closest, so dropping it is always safe. After the last point, the heap holds exactly the k closest.

    Each push and pop costs O(log k), because the heap never holds more than k + 1 entries, and there are n pushes, so the scan takes O(n log k).

    The heap entries are array<int, 3> values {dist, x, y}. A std::array compares element by element, so entries are ordered by distance first, then x, then y, which is exactly the order the answer needs. For {{3, 3}, {5, -1}, {-2, 4}} with k = 2 the scan goes like this:

    after push (3,3): size 1 top dist 18
after push (5,-1): size 2 top dist 26
after push (-2,4): size 3 top dist 26 -> pop dist 26

    The third push makes the heap too big, and the top, the point at squared distance 26, is dropped. The heap is left with the points at 18 and 20.

    A heap is only partly ordered, so the kept entries still have to be sorted. They come out of the heap farthest first, so the code collects them into a vector and sorts it. That sort handles only k entries, so it costs O(k log k), which never exceeds O(n log k) because k <= n.

    Step by step

    1. Declare the heap: priority_queue<array<int, 3>> heap, a max-heap of {dist, x, y} entries.
    2. Visit every point with a range-based loop, for (auto& point : points). The & reads each point in place instead of copying it.
    3. Read the coordinates into x and y.
    4. Compute the squared distance dist as x * x + y * y.
    5. Push {dist, x, y}. The braces build the array<int, 3> in the call.
    6. Trim the heap: if heap.size() > k, call heap.pop(). After every point the heap holds the k closest seen so far.
    7. Declare closest, a vector<array<int, 3>>, to take the entries out. A priority_queue can only be read from its top.
    8. Loop while !heap.empty(), which runs k times.
    9. Move each entry out: push_back(heap.top()), then heap.pop(). In C++ pop() returns nothing, so the top is read first and removed second.
    10. Sort closest from begin() to end(). The default order on arrays is distance, then x, then y.
    11. Declare result, a vector<vector<int>>.
    12. Loop over closest with auto& entry.
    13. Convert each entry into a point with result.push_back({entry[1], entry[2]}). The distance at index 0 has done its job.
    14. Return result.

    The solution in C++

    class Solution {
public:
    vector<vector<int>> kClosest(vector<vector<int>>& points, int k) {
        priority_queue<array<int, 3>> heap;
        for (auto& point : points) {
            int x = point[0], y = point[1];
            int dist = x * x + y * y;
            heap.push({dist, x, y});
            if (heap.size() > k) {
                heap.pop();
            }
        }
        vector<array<int, 3>> closest;
        while (!heap.empty()) {
            closest.push_back(heap.top());
            heap.pop();
        }
        sort(closest.begin(), closest.end());
        vector<vector<int>> result;
        for (auto& entry : closest) {
            result.push_back({entry[1], entry[2]});
        }
        return result;
    }
};

    The C++ details worth a closer look:

    • priority_queue is a max-heap. With no comparator argument it keeps the largest element on top, which is what this problem needs. Finding the smallest values with a heap of size k uses the largest of them as the gatekeeper.
    • array<int, 3> compares lexicographically. operator< compares the first elements, then the second if they are equal, then the third. A heap or sort over these arrays needs no custom comparator.
    • pop() returns nothing. Read top() first, then call pop(). The code copies the entry into closest before removing it.
    • heap.size() > k compares an unsigned size with an int. The constraints keep k at 1 or more, so the comparison is correct, though some compilers warn about the mix.
    • No overflow. x * x + y * y is at most 200,000,000, well inside the int range, so the squared distance never needs long long here. Signed overflow would be undefined behaviour.
    • Brace initialisation. heap.push({dist, x, y}) and result.push_back({entry[1], entry[2]}) build an array<int, 3> and a vector<int> directly in the call.

    Complexity

    Time: O(n log k). The loop over the points does one push and at most one pop per point, each O(log k) because the heap never holds more than k + 1 entries. Emptying the heap costs O(k log k), and sorting the k entries costs O(k log k). Since k <= n, the total is O(n log k).

    Space: O(k) extra. The heap holds at most k + 1 entries, and closest and result hold k each.

    Edge cases and common mistakes

    • Using a min-heap. A heap whose top is the closest point would need all n points inside it. Keeping the farthest point on top is what allows the heap to stay at size k.
    • Skipping the final sort. The heap pops points farthest first, and its internal order isn't sorted anyway. The answer needs the closest first, in the order given by the statement.
    • Taking square roots. Square roots add nothing and bring in floating-point rounding. Comparing squared distances is exact.
    • Losing the tie-break. Ordering entries by distance alone gives no rule for equal distances. Keeping x and y in the entry makes the tie-break part of the comparison. With the points {0, 2}, {2, 0} and {0, -2}, all at squared distance 4, the result is [0,-2], [0,2], [2,0].
    • Duplicates. The same point can appear more than once. {{1, 1}, {1, 1}, {2, 2}} with k = 2 gives [[1,1],[1,1]].
    • k equal to the number of points. Nothing is ever popped. {{3, 3}, {5, -1}, {-2, 4}} with k = 3 gives [[3,3],[-2,4],[5,-1]], all the points sorted.
    • Coordinates at the limits. {{10000, 10000}, {-10000, -10000}} with k = 2 gives [[-10000,-10000],[10000,10000]]: equal distances, so the smaller x comes first.

    Testing

    • points = {{1, 3}, {-2, 2}}, k = 1 gives [[-2,2]].
    • points = {{3, 3}, {5, -1}, {-2, 4}}, k = 2 gives [[3,3],[-2,4]].
    • points = {{0, 2}, {2, 0}, {1, 1}, {0, -2}, {3, 3}}, k = 4 gives [[1,1],[0,-2],[0,2],[2,0]].
    • points = {{5, 5}}, k = 1 gives [[5,5]].
    • points = {{1, 1}, {1, 1}, {2, 2}}, k = 2 gives [[1,1],[1,1]].
    Keep going with C++500+ step-by-step lessons, from your first line of code to coding interviews. Free to start.
    Start learning free