LeetCode 973: K Closest Points to Origin

Problem Restatement

We are given an array points, where each item is a point on the 2D plane:

[x, y]

We are also given an integer k.

Return the k points closest to the origin:

(0, 0)

The distance is Euclidean distance. The answer may be returned in any order, and the answer is guaranteed to be unique except for order. The official constraints are 1 <= k <= points.length <= 10^4 and -10^4 <= xi, yi <= 10^4.

Input and Output

Example function shape:

def kClosest(points: list[list[int]], k: int) -> list[list[int]]:
    ...

Examples

Example 1:

points = [[1, 3], [-2, 2]]
k = 1

Output:

[[-2, 2]]

Explanation:

[1, 3]   has distance sqrt(10)
[-2, 2]  has distance sqrt(8)

Since sqrt(8) < sqrt(10), the closest point is [-2, 2].

Example 2:

points = [[3, 3], [5, -1], [-2, 4]]
k = 2

Output:

[[3, 3], [-2, 4]]

The output may also be:

[[-2, 4], [3, 3]]

because any order is accepted.

First Thought: Compute Distances and Sort

The most direct method is:

1. Compute the distance of every point from the origin.
2. Sort the points by distance.
3. Return the first k points.

We do not need the actual square root.

For a point:

[x, y]

the Euclidean distance is:

sqrt(x * x + y * y)

Since square root is increasing, comparing:

x * x + y * y

gives the same ordering.

So we sort by squared distance.

Key Insight

The square root does not change which distance is smaller.

If:

a < b

then:

sqrt(a) < sqrt(b)

So instead of sorting by real distance, we sort by squared distance:

x * x + y * y

This avoids floating-point arithmetic and keeps the implementation simple.

Algorithm

1. Sort points using squared distance from the origin as the key.
2. Return the first k points.

The sorting key is:

point[0] * point[0] + point[1] * point[1]

Correctness

For every point [x, y], its distance from the origin is determined by x * x + y * y.

Because square root is strictly increasing for nonnegative numbers, sorting by squared distance produces the same order as sorting by Euclidean distance.

After sorting, the first k points are exactly the k smallest distances from the origin.

The problem accepts the answer in any order, so returning these first k sorted points is valid.

Therefore, the algorithm returns exactly the required closest points.

Complexity

Let n = len(points).

Implementation

class Solution:
    def kClosest(self, points: list[list[int]], k: int) -> list[list[int]]:
        points.sort(key=lambda point: point[0] * point[0] + point[1] * point[1])
        return points[:k]

Code Explanation

We sort the points in place:

points.sort(key=lambda point: point[0] * point[0] + point[1] * point[1])

For each point, the key is its squared distance from the origin.

Then we return the first k points:

return points[:k]

These are the closest points after sorting.

Testing

Because output order may vary, compare sorted normalized results.

def normalize(points):
    return sorted(map(tuple, points))

def run_tests():
    s = Solution()

    result = s.kClosest([[1, 3], [-2, 2]], 1)
    assert normalize(result) == normalize([[-2, 2]])

    result = s.kClosest([[3, 3], [5, -1], [-2, 4]], 2)
    assert normalize(result) == normalize([[3, 3], [-2, 4]])

    result = s.kClosest([[0, 0], [1, 1], [2, 2]], 1)
    assert normalize(result) == normalize([[0, 0]])

    result = s.kClosest([[1, 0], [0, 1], [2, 2]], 2)
    assert normalize(result) == normalize([[1, 0], [0, 1]])

    result = s.kClosest([[-5, -4], [3, 2], [1, 1]], 2)
    assert normalize(result) == normalize([[3, 2], [1, 1]])

    print("all tests passed")

run_tests()

Test meaning: