Closest Points To Origin
Problem
Given points on a plane and a number k, return the k points lying closest to the origin. Distance is the ordinary straight-line distance. Report the chosen points ordered by distance, and when two points are equally far away prefer the one with the smaller first coordinate, then the smaller second one.
Examples
Input: points = [[3, 3], [5, -1], [-2, 4]], k = 2
Output: [(3, 3), (-2, 4)]
Why: their squared distances are 18 and 20, while the middle point sits at 26Input: points = [[1, 0], [0, 1]], k = 1
Output: [(0, 1)]
Why: edge case, an exact tie is broken by the smaller first coordinateInput: points = [[1, 3], [-2, 2]], k = 1
Output: [(-2, 2)]Hints
0 / 3
Comparing distances does not need the square root at all: squaring is an increasing function, so the squared distance ranks the points identically and avoids floating point.
Sorting everything answers a question you did not ask, since only k points matter. Think about holding exactly k candidates and evicting the worst one whenever a better point arrives.
Keep a structure of at most k points that always exposes the furthest one held. Add each point, and when the collection grows past k, remove that furthest point. Whatever survives to the end is the answer, which you then order for reporting.
Solution
Ranking by squared distance avoids the square root entirely without changing the order. Holding at most k candidates in a heap whose top is the furthest one means each new point costs a logarithm of k rather than a full sort, and the eviction rule keeps exactly the best k. Negating the key turns the library's min-heap into the max-heap this needs, and negating the coordinates makes the eviction prefer to drop the larger ones on a tie. Time is O(n log k) and space is O(k).
import heapq
def closest_points(points, k):
heap = [] # at most k points, the furthest one on top
for x, y in points:
heapq.heappush(heap, (-(x * x + y * y), -x, -y))
if len(heap) > k:
heapq.heappop(heap) # drop the furthest point held so far
chosen = [(-a, -b) for _, a, b in heap]
return sorted(chosen, key=lambda p: (p[0] * p[0] + p[1] * p[1], p))
print(closest_points([[3, 3], [5, -1], [-2, 4]], 2)) # -> [(3, 3), (-2, 4)]
print(closest_points([[1, 0], [0, 1]], 1)) # -> [(0, 1)]
print(closest_points([[1, 3], [-2, 2]], 1)) # -> [(-2, 2)]Stuck on the idea rather than the code? K Closest Points covers it.