K Closest Points to Origin
You get n points on a plane: point i is (xs[i], ys[i]). Return the k points closest to the origin (0, 0), each as a pair [x, y].
Order the answer by distance from the origin, nearest first. Points at the same distance go by smaller x first, then by smaller y. A point listed more than once counts once per copy.
Function
- xsinteger-array
- the x coordinate of each point
- ysinteger-array
- the y coordinate of each point
- kinteger
- how many points to return
- Returnsinteger-2d-array
- the k closest points as [x, y] pairs, nearest first, ties by x and then y
Constraints
n == xs.length == ys.length1 ≤ k ≤ n ≤ 5000-104 ≤ xs[i], ys[i] ≤ 104- Points may repeat; each copy is a separate point.
Examples
- Input
- xs = [3, -1, 2, 0]ys = [1, 2, -2, 5]k = 2
- Output
- [[-1, 2], [2, -2]]
- Explanation
- The squared distances are 10, 5, 8 and 25. The two smallest belong to
(-1, 2)with 5 and(2, -2)with 8, in that order.
- Input
- xs = [1, -1, 0, 1]ys = [1, 1, 2, -1]k = 3
- Output
- [[-1, 1], [1, -1], [1, 1]]
- Explanation
(1, 1),(-1, 1)and(1, -1)all have squared distance 2, and(0, 2)has 4. Among the tied points(-1, 1)has the smallest x, and(1, -1)comes before(1, 1)by y.
- Input
- xs = [7]ys = [-7]k = 1
- Output
- [[7, -7]]
- Explanation
- There is only one point and you need one, so the answer is
(7, -7)itself.
Follow-up
The points now arrive one at a time, and after each one you must report the k closest so far. How much memory do you need?
Hints
Open them one at a time. Each one gives away a little more.
You never need the actual distance. What can you compare instead of
√(x² + y²)and still get the same order?Say you keep the
kclosest points seen so far. Which one of them does a new point compete with?Keep a max-heap of at most
kpoints ordered by(x² + y², x, y). Replace the top when a new point comes before it, and at the end take the points off the heap to fill the answer from the back.
Solution
Distance needs no square root: x² + y² gives the same order and stays an integer. You need only k of the n points, so a max-heap of the k closest so far, with the farthest on top ready to be replaced, costs O(n log k) instead of a full sort.
Pick the closest point k times
Correct, but does not finish on the largest tests
Intuition
The first point of the answer is the closest of all, the second is the closest of the rest, and so on. So scan the points k times, each time taking the point that comes first by (x² + y², x, y) among those not taken yet.
Each scan touches all n points. For n = 5000 and k = 2500 that is 1.25 × 10^7 three-part comparisons.
Algorithm
- Mark every point as not taken.
- Repeat
ktimes: scan the points not taken and find the one that comes first by(x² + y², x, y). - Mark it taken and append
[x, y]to the answer. - Return the answer.
def kClosest(xs, ys, k):
n = len(xs)
taken = [False] * n
result = []
for _ in range(k):
best = -1
for i in range(n):
# Compare (squared distance, x, y): the order the answer uses.
if not taken[i] and (best == -1 or
(xs[i] * xs[i] + ys[i] * ys[i], xs[i], ys[i]) <
(xs[best] * xs[best] + ys[best] * ys[best], xs[best], ys[best])):
best = i
taken[best] = True
result.append([xs[best], ys[best]])
return resultSort every point
Intuition
The answer is the first k points of the full order, so sort every point by (x² + y², x, y) and keep the first k. In example 1 the sorted keys are (5, -1, 2), (8, 2, -2), (10, 3, 1) and (25, 0, 5), and the first two give the answer.
The sort costs O(n log n) whatever k is. When k is small, most of that work orders points that never reach the answer.
Algorithm
- Build each point's key
(x² + y², x, y). - Sort the points by key.
- Return the first
kof them as[x, y]pairs.
def kClosest(xs, ys, k):
points = list(zip(xs, ys))
# Sort by squared distance, then x, then y: the order the answer uses.
points.sort(key=lambda p: (p[0] * p[0] + p[1] * p[1], p[0], p[1]))
return [[x, y] for x, y in points[:k]]Keep the k closest in a max-heap
Intuition
Keep the best k points seen so far. A newcomer can only push out the kept point that ranks last, so store the kept points in a max-heap by (x² + y², x, y), with that point on top. Once k points are stored, a point that ranks ahead of the top replaces it; any other point is dropped, because k stored points already rank ahead of it.
In example 2 with k = 3, the top after the first three points is (0, 2) with key (4, 0, 2). The point (1, -1) has key (2, 1, -1), ranks ahead, and replaces it.
A heap is only partly sorted, so remove the top k times and fill the answer from its last slot to its first. Each point costs at most one O(log k) heap step, so the time is O(n log k) and the memory O(k).
Algorithm
- Make an empty max-heap keyed by
(x² + y², x, y). - Insert each of the first
kpoints. - For every later point, compare it with the top: if it ranks ahead, swap it in and sift it down; otherwise skip it.
- Remove the top
ktimes; each removed point fills the last free slot of the answer. - Return the answer.
import heapq
def kClosest(xs, ys, k):
# heapq is a min-heap; negating every part of (distance, x, y) puts the
# farthest kept point, in the answer's order, on top.
heap = []
for x, y in zip(xs, ys):
item = (-(x * x + y * y), -x, -y)
if len(heap) < k:
heapq.heappush(heap, item)
elif item > heap[0]:
# The new point comes before the farthest kept one: swap them.
heapq.heapreplace(heap, item)
# Popping gives the farthest first, so fill the answer from the back.
result = [None] * len(heap)
for i in range(len(result) - 1, -1, -1):
_, x, y = heapq.heappop(heap)
result[i] = [-x, -y]
return result
Pitfalls and edge cases
- Skipping the tie-break. Example 2 has three points at distance 2; without comparing
xand thenythey come out in input order instead of(-1, 1),(1, -1),(1, 1). - Keeping a min-heap of
kpoints. Its top is the nearest point, so each eviction throws away a point the answer needs, and you end with thekfarthest. - Reading the answer straight out of the heap's array. That array is only partly ordered; remove the points one by one first.
- Removing repeated points. With
(5, 5)listed three times andk = 3, the answer is that point three times.
Frequently asked questions4
What is the time complexity of K Closest Points to Origin?
Sorting all points takes O(n log n). The heap version is O(n log k) for the scan plus O(k log k) to order the k survivors, and it keeps only O(k) points in memory.
Why use a max-heap, not a min-heap, for the k closest points?
A new point competes with the farthest of the k kept points, so that point must be on top. A min-heap works only if it holds all n points so you can remove the k nearest, which costs O(n) memory instead of O(k).
Do you need the square root to compare distances?
No. For numbers that are not negative, a < b exactly when √a < √b, so x² + y² gives the same order as the distance. It also stays an integer: at most 2 × 10^8 here, inside the 32-bit range.
Can K Closest Points to Origin be solved in O(n) time?
On average, yes. Quickselect partitions the points around a random pivot by the same (x² + y², x, y) order and continues only on the side that holds position k, in O(n) on average. The k points left of that position are the answer; sorting them adds O(k log k).
Similar problems
Problems that use the same ideas. Solving two or three of them is what makes a pattern stick.
def kClosest(xs, ys, k):
# Write code hereTest #1
Test #2
Test #3
Arguments
xs = [3, -1, 2, 0] ys = [1, 2, -2, 5] k = 2
Expected Output
[[-1, 2], [2, -2]]