K Closest Points to Origin

Difficulty: Medium Source: NeetCode

Problem

Given an array of points where points[i] = [xi, yi] represents a point on the X-Y plane, and an integer k, return the k closest points to the origin (0, 0).

The distance between two points on the X-Y plane is the Euclidean distance: sqrt((x1 - x2)² + (y1 - y2)²).

You may return the answer in any order. The answer is guaranteed to be unique (except for the order of the answer).

Example 1: Input: points = [[1, 3], [-2, 2]], k = 1 Output: [[-2, 2]]

Example 2: Input: points = [[3, 3], [5, -1], [-2, 4]], k = 2 Output: [[3, 3], [-2, 4]]

Constraints:

• 1 <= k <= points.length <= 10^4
• -10^4 <= xi, yi <= 10^4

Prerequisites

Before attempting this problem, you should be comfortable with:

• Euclidean distance — you can skip the sqrt when comparing distances since sqrt is monotonic
• Heaps / Priority Queues — maintaining a bounded collection of k elements efficiently
• Python’s heapq module — heappush, heappop, and nsmallest

1. Brute Force — Sort by Distance

Intuition

Compute the squared distance for every point, sort by that distance, and slice off the first k results. Simple and correct — the only downside is that we sort all n points even though we only need k of them.

Algorithm

1. For each point, compute x² + y² (no need for sqrt since we’re just comparing).
2. Sort the points by this distance.
3. Return the first k points.

Solution

def kClosest(points, k):
    points.sort(key=lambda p: p[0]**2 + p[1]**2)
    return points[:k]


print(kClosest([[1, 3], [-2, 2]], 1))            # [[-2, 2]]
print(kClosest([[3, 3], [5, -1], [-2, 4]], 2))   # [[3, 3], [-2, 4]]
print(kClosest([[0, 1], [1, 0]], 2))             # [[0, 1], [1, 0]]

Complexity

• Time: O(n log n)
• Space: O(1) extra (sort is in-place)

2. Max-Heap of Size K

Intuition

We want to keep exactly k closest points. Use a max-heap (by distance) of size k. For every new point, push it onto the heap. If the heap grows beyond k, pop the farthest point — it’s no longer in the top-k closest. When all points are processed, everything left in the heap is a winner.

Since Python’s heapq is a min-heap, store (-distance, x, y) so the largest distance sits at the root and gets popped first.

Algorithm

1. Initialize an empty heap.
2. For each point [x, y]: a. Push (-dist, x, y) onto the heap where dist = x² + y². b. If len(heap) > k, pop the root (farthest point).
3. Return [[x, y] for _, x, y in heap].

Solution

import heapq

def kClosest(points, k):
    heap = []  # max-heap via negation

    for x, y in points:
        dist = x * x + y * y
        heapq.heappush(heap, (-dist, x, y))
        if len(heap) > k:
            heapq.heappop(heap)  # remove farthest

    return [[x, y] for _, x, y in heap]


print(kClosest([[1, 3], [-2, 2]], 1))            # [[-2, 2]]
print(kClosest([[3, 3], [5, -1], [-2, 4]], 2))   # [[3, 3], [-2, 4]]
print(kClosest([[0, 1], [1, 0]], 2))             # [[0, 1], [1, 0]]

Complexity

• Time: O(n log k) — each push/pop on a heap of size k costs O(log k)
• Space: O(k)

3. Bonus — heapq.nsmallest

Intuition

Python’s standard library has heapq.nsmallest(k, iterable, key) which does exactly this under the hood. Worth knowing for interviews and quick scripting.

Solution

import heapq

def kClosest(points, k):
    return heapq.nsmallest(k, points, key=lambda p: p[0]**2 + p[1]**2)


print(kClosest([[1, 3], [-2, 2]], 1))            # [[-2, 2]]
print(kClosest([[3, 3], [5, -1], [-2, 4]], 2))   # [[3, 3], [-2, 4]]

Complexity

• Time: O(n log k) internally
• Space: O(k)

Common Pitfalls

Taking the square root unnecessarily. Since sqrt is a monotonically increasing function, sqrt(a) < sqrt(b) iff a < b. Skip the sqrt — it saves computation and avoids floating-point precision issues.

Confusing min-heap and max-heap semantics. We want to evict the farthest point when the heap overflows, so we need the farthest at the top. That means a max-heap (negate distances) even though we ultimately want the closest.

Returning distances instead of points. The problem asks for the points themselves, not their distances. Make sure the heap stores the coordinates, not just the distance value.