Start learning free

LESSON 385 OF 536 · HEAP

K Closest Points to Origin in Python: O(n log k) with a Heap

Practice O(n log k) Read the full explanation

Given a list 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 8
    Ctrl + Enter to check

    PYTHON · HEAP

    K Closest Points To Origin in Python, explained

    Problem overview

    The method kClosest receives points, a list where each point is [x, y], and an integer k. It returns the k points closest to the origin (0, 0), as a list of [x, y] lists.

    The result is sorted: closest first, points at the same distance by x, and points with the same x by y. The answer is unique: every point left out is strictly farther than every point kept.

    For points = [[3, 3], [5, -1], [-2, 4]] and k = 2, the squared distances are 18, 26 and 20. The two closest are [3, 3] and [-2, 4], in that order, so the result is [[3, 3], [-2, 4]]. For points = [[1, 3], [-2, 2]] and k = 1, the point [-2, 2] is closer, so the result is [[-2, 2]].

    In a third example, points = [[0, 2], [2, 0], [1, 1], [0, -2], [3, 3]] and k = 4, three points are at the same distance, 2. They are ordered by x and then by y, so the result is [[1, 1], [0, -2], [0, 2], [2, 0]].

    Approach: keep only the k best so far

    The simplest solution sorts all the points by distance and takes the first k, which costs O(n log n). The statement asks for O(n log k) time and O(k) extra space, which means never storing more than k points at a time.

    Two observations make that possible.

    • No square root is needed. A point is closer than another exactly when its x * x + y * y is smaller. The squared distance stays a whole number, so there is no float rounding to worry about.
    • The point to throw away is the farthest one. Keep a collection of the best k points. When a new point arrives, add it. If the collection now has k + 1 points, the farthest of them can't be among the k closest, so remove it.

    The structure that finds and removes the farthest element quickly is a heap whose top is the largest distance. Python's heapq only provides a min-heap, where the smallest value is on top. The usual trick is to store the distance negated: the smallest negated distance is the largest distance, so the farthest point kept sits on top. Each push and pop costs O(log k), and there are n pushes, so the total is O(n log k).

    Here is the heap on points = [[3, 3], [5, -1], [-2, 4]] with k = 2. Each entry is (-dist, x, y):

    pushed (3, 3) [(-18, 3, 3)] 1
pushed (5, -1) [(-26, 5, -1), (-18, 3, 3)] 2
pushed (-2, 4) [(-26, 5, -1), (-20, -2, 4), (-18, 3, 3)] 3
popped (-26, 5, -1)
heap now [(-20, -2, 4), (-18, 3, 3)]

    After the third push the heap holds three entries, more than k, so the top, (-26, 5, -1), is removed. That is the point at squared distance 26, the farthest. (The lists above are sorted for display; inside the heap, only the top is guaranteed to be the smallest.)

    Step by step

    1. Start with an empty list named heap. Using heapq functions on it turns it into a heap that holds the k closest points seen so far.
    2. Visit every point with for x, y in points. Each point is a list of two numbers, so it unpacks straight into two names.
    3. Compute the squared distance dist = x * x + y * y.
    4. Push (-dist, x, y) with heapq.heappush. Tuples compare by their first item first, so the heap orders entries by -dist, with the farthest point on top.
    5. Trim to size. If len(heap) > k, call heapq.heappop(heap). Out of k + 1 points, the farthest one can't be in the answer.
    6. Turn the entries back into positive distances with a list comprehension, closest, that builds (-negDist, x, y) from each entry. A heap is only partly ordered, so the points still need sorting.
    7. Sort closest in place with closest.sort(). Tuples compare item by item: distance, then x, then y, which is exactly the order the answer needs.
    8. Return the coordinates only, with a list comprehension that builds [x, y] from each dist, x, y.

    The solution in Python

    class Solution:
    def kClosest(self, points: List[List[int]], k: int) -> List[List[int]]:
        heap = []
        for x, y in points:
            dist = x * x + y * y
            heapq.heappush(heap, (-dist, x, y))
            if len(heap) > k:
                heapq.heappop(heap)
        closest = [(-negDist, x, y) for negDist, x, y in heap]
        closest.sort()
        return [[x, y] for dist, x, y in closest]

    Python details worth a closer look:

    • heapq is a min-heap on a plain list. There is no heap class: heappush and heappop keep an ordinary list arranged so that heap[0] is the smallest item. Printing the list shows it in heap order, which is not sorted order.
    • Tuples compare item by item. (-18, 3, 3) < (-20, -2, 4) is False, because -18 is larger than -20. If the first items are equal, the second decide, then the third. This is what lets the heap carry x and y along with the distance, and it is also what puts equal distances in order of x and y when closest is sorted.
    • Negating for a max-heap. Pushing -dist turns "smallest on top" into "farthest on top". The comprehension that builds closest negates again to restore the real distances before sorting.
    • sort() works in place. It reorders closest and returns None, so return closest.sort() would return None. The solution sorts on one line and builds the answer on the next.
    • Unpacking in comprehensions. for negDist, x, y in heap unpacks each three-item tuple, and for dist, x, y in closest does the same. The name dist in the last line is not used; it only has to take the first item.
    • No overflow. The squared distance can reach 2 * 10,000 * 10,000, and Python integers can't overflow anyway.

    Complexity

    Time: O(n log k). Each of the n points costs one push and at most one pop on a heap of at most k + 1 entries, so each is O(log k). The final sort handles k entries, which is O(k log k), and that is never more than O(n log k) because k <= n.

    Space: O(k). The heap holds at most k + 1 entries, and closest holds k. The returned list of k points is output.

    Edge cases and common mistakes

    • k equals the number of points. Nothing is ever popped, and the method returns every point, sorted. For [[3, 3], [5, -1], [-2, 4]] with k = 3 the result is [[3, 3], [-2, 4], [5, -1]].
    • A point at the origin. [[0, 0]] with k = 1 gives [[0, 0]]. The distance is 0, and 0 negated is still 0.
    • Duplicate points. [[1, 1], [1, 1], [2, 2]] with k = 2 gives [[1, 1], [1, 1]]. Duplicates are separate entries and each counts toward k.
    • Ties on the distance. The tuples carry x and y, so equal distances are ordered consistently. The statement promises the answer is unique, so the boundary between kept and dropped points never falls inside a tie.
    • Forgetting to negate. Pushing dist instead of -dist keeps the closest point on top, so the pops remove the closest points and the method returns the k farthest.
    • Returning the heap's order. A heap is not sorted. Without the final sort, the points come out in heap order, which the lesson's tests reject.
    • Using the square root. It isn't needed. It also adds float values to the comparison for no benefit.

    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 = [[3, 3], [5, -1], [-2, 4]], k = 3 gives [[3, 3], [-2, 4], [5, -1]].
    • points = [[0, 0]], k = 1 gives [[0, 0]].
    • points = [[1, 1], [1, 1], [2, 2]], k = 2 gives [[1, 1], [1, 1]].
    Keep going with Python500+ step-by-step lessons, from your first line of code to coding interviews. Free to start.
    Start learning free