K Closest Points to Origin

Medium Top 250
Associated Patterns
Interviewed At (Company Tags)
AmazonGoogleMetaMicrosoftUber

Problem Description

Given an array of points where points[i] = [xᵢ, yᵢ] 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 is the Euclidean distance: √(x² + y²).

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


Examples

Example 1:

Input: points = [[1,3],[-2,2]], k = 1 Output: [[-2,2]] Explanation: The distances are √10 and √8. The closest point is [-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⁴
  • -10⁴ ≤ xᵢ, yᵢ ≤ 10⁴

Why Sorting Is Wasteful for Large Inputs

Sorting all points takes O(n log n) and computes the full order, which you don’t need. You only need the k smallest distances. A max-heap of size k gives you exactly that in O(n log k) time.

Keep a max-heap with the k closest points seen so far. For each new point, compare its squared distance to the farthest point in the heap (the top). If it’s closer, pop the farthest and push the new one. At the end, the heap contains exactly the k closest points.

Skip the square root. Comparing squared distances preserves the correct order and avoids floating-point overhead.


Solution 1: Max-Heap of Size k

import java.util.*;

class Solution {
    public int[][] kClosest(int[][] points, int k) {
        // max-heap: compare by squared distance, largest on top
        PriorityQueue<int[]> heap = new PriorityQueue<>(
            (a, b) -> (b[0]*b[0] + b[1]*b[1]) - (a[0]*a[0] + a[1]*a[1])
        );

        for (int[] point : points) {
            heap.offer(point);
            if (heap.size() > k) {
                heap.poll(); // remove the farthest point
            }
        }

        return heap.toArray(new int[k][]);
    }
}
import heapq

class Solution:
    def kClosest(self, points: list[list[int]], k: int) -> list[list[int]]:
        # Python's heapq is a min-heap; negate distance for max-heap behavior
        heap = []

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

        return [[x, y] for _, x, y in heap]
#include <vector>
#include <queue>

class Solution {
public:
    std::vector<std::vector<int>> kClosest(std::vector<std::vector<int>>& points, int k) {
        // max-heap: sort by squared distance descending
        auto cmp = [](const std::vector<int>& a, const std::vector<int>& b) {
            return a[0]*a[0] + a[1]*a[1] < b[0]*b[0] + b[1]*b[1];
        };
        std::priority_queue<std::vector<int>, std::vector<std::vector<int>>, decltype(cmp)> heap(cmp);

        for (auto& point : points) {
            heap.push(point);
            if ((int)heap.size() > k) heap.pop(); // remove farthest
        }

        std::vector<std::vector<int>> result;
        while (!heap.empty()) {
            result.push_back(heap.top());
            heap.pop();
        }
        return result;
    }
};

Complexity Analysis:

  • Time Complexity: O(n log k). Each of the n points is pushed and potentially popped from a heap of size k.
  • Space Complexity: O(k) for the heap.

Where it breaks: if k is close to n, the heap advantage over sorting diminishes. At k = n, sorting is equivalent and potentially faster in practice due to cache efficiency.


Solution 2: Sort (Simpler, Slightly Slower)

import java.util.*;

class Solution {
    public int[][] kClosest(int[][] points, int k) {
        Arrays.sort(points, (a, b) -> (a[0]*a[0] + a[1]*a[1]) - (b[0]*b[0] + b[1]*b[1]));
        return Arrays.copyOf(points, k);
    }
}
class Solution:
    def kClosest(self, points: list[list[int]], k: int) -> list[list[int]]:
        points.sort(key=lambda p: p[0]**2 + p[1]**2)
        return points[:k]
#include <vector>
#include <algorithm>

class Solution {
public:
    std::vector<std::vector<int>> kClosest(std::vector<std::vector<int>>& points, int k) {
        std::sort(points.begin(), points.end(), [](const auto& a, const auto& b) {
            return a[0]*a[0] + a[1]*a[1] < b[0]*b[0] + b[1]*b[1];
        });
        return std::vector<std::vector<int>>(points.begin(), points.begin() + k);
    }
};

Complexity Analysis:

  • Time Complexity: O(n log n).
  • Space Complexity: O(log n) for the sort stack.

Where it breaks: when the stream of points arrives one at a time (online), sorting isn’t possible. The heap handles streaming inputs naturally.


Common Mistakes

  • Taking the square root for distance comparison. Comparing x²+y² directly is equivalent and avoids floating-point imprecision and the overhead of sqrt.
  • Using a min-heap of size n and popping k times. This gives the right answer but uses O(n) space and O(n log n) time, eliminating the heap’s advantage.

Frequently Asked Questions

Can this be solved in O(n) average time? Yes, using QuickSelect (partial sort). It finds the k-th smallest element in O(n) average, O(n²) worst case. For this problem’s constraints, the heap or sort approach is simpler and fast enough.

What if the input is a stream and k is fixed? Use the max-heap approach: process each incoming point and maintain a heap of size k. When a closer point arrives, evict the current farthest.


← All Problems