66.K Closest Points to Origin

MediumHeap / Priority Queue · 3/7

Problem

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

Distance is Euclidean, sqrt(x² + y²). The answer may be returned in any order and is guaranteed to be unique (apart from order).

Examples

Input:  points = [[1,3],[-2,2]], k = 1
Output: [[-2,2]]
Explanation: dist²(1,3) = 10, dist²(-2,2) = 8. 8 < 10, so [-2,2] is closer.

Input:  points = [[3,3],[5,-1],[-2,4]], k = 2
Output: [[3,3],[-2,4]]
Explanation: squared distances are 18, 26, 20 -> keep 18 and 20.

Constraints

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

💡 Key Insight

Compare squared distances x² + y² — the square root is monotonic, so it never changes the ordering and avoids floating point. We do not need the points fully sorted, only partitioned into "k closest" and "the rest": a size-k max-heap does that in O(n log k), and Quickselect does it in O(n) on average.

Brute Force (Sort by Distance)

TimeO(n log n)SpaceO(log n) to O(n) for the sort, plus O(k) for the output.

Intuition

Sort all points by squared distance and take the first k.

Algorithm

Sort points with comparator x² + y².
Return Arrays.copyOfRange(points, 0, k).

Visualization

points:  [3,3] d=18   [5,-1] d=26   [-2,4] d=20
sorted:  [3,3] 18, [-2,4] 20, [5,-1] 26
k = 2 -> [[3,3],[-2,4]]

Code

java

import java.util.Arrays;

class Solution {
    public int[][] kClosest(int[][] points, int k) {
        // squared distance is enough for ordering
        Arrays.sort(points, (a, b) -> Integer.compare(a[0] * a[0] + a[1] * a[1], b[0] * b[0] + b[1] * b[1]));
        return Arrays.copyOfRange(points, 0, k);
    }
}

Complexity

Time: O(n log n) — sorting.
Space: O(log n) to O(n) for the sort, plus O(k) for the output.

Max-Heap of Size k

TimeO(n log k)SpaceO(k)

Intuition

Scan the points, keeping the k closest seen so far in a max-heap keyed by distance. The root is the farthest of those k; whenever a new point pushes the heap above k, evict the root. At the end the heap holds exactly the k closest.

Algorithm

Create a max-heap ordered by squared distance (descending).
For each point: offer it; if size() > k, poll().
Drain the heap into the result array.

Visualization

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

flowchart LR
    A["offer [3,3] d=18<br/>heap = {18}"] --> B["offer [5,-1] d=26<br/>heap = {26,18}"]
    B --> C["offer [-2,4] d=20<br/>heap = {26,20,18}"]
    C -->|"size 3 > k, poll 26"| D["heap = {20,18}"]
    D --> E["answer [[-2,4],[3,3]]"]

Code

java

import java.util.PriorityQueue;

class Solution {
    public int[][] kClosest(int[][] points, int k) {
        // max-heap: farthest point on top
        PriorityQueue<int[]> maxHeap = new PriorityQueue<>((a, b) -> Integer.compare(dist(b), dist(a)));
        for (int[] p : points) {
            maxHeap.offer(p);
            if (maxHeap.size() > k) {
                maxHeap.poll(); // evict the farthest of k + 1
            }
        }
        int[][] result = new int[k][];
        for (int i = 0; i < k; i++) {
            result[i] = maxHeap.poll();
        }
        return result;
    }

    private int dist(int[] p) {
        return p[0] * p[0] + p[1] * p[1];
    }
}

Complexity

Time: O(n log k) — each of n points costs O(log k) heap work.
Space: O(k) — heap size.

QuickselectOptimal

TimeO(n) averageSpaceO(1) extra (iterative), plus O(k) for the output.

Intuition

Quickselect partitions the array around a random pivot distance: closer points move left, farther points move right, and the pivot lands at its final sorted index p. If p == k, the first k slots are exactly the k closest. Otherwise recurse (iteratively) into only the side containing index k. Discarding one side each round gives O(n) expected time.

Algorithm

lo = 0, hi = n - 1.
While lo < hi: p = partition(lo, hi) with a random pivot.
- If p == k, stop. If p < k, set lo = p + 1; else set hi = p - 1.
Return the first k points.
partition (Lomuto): move a random pivot to hi, sweep i from lo to hi - 1, swapping points with smaller distance to the front, then place the pivot at store.

Dry Run

k = 2, distances by index: [18, 26, 20]   (points [3,3], [5,-1], [-2,4])

partition(0, 2), pivot = index 2 (d = 20):
  i=0: 18 < 20 -> swap into store 0, store = 1
  i=1: 26 >= 20 -> skip
  place pivot at store 1 -> distances [18, 20, 26], p = 1
p = 1 < k = 2 -> lo = 2; lo == hi -> stop
first k = [[3,3],[-2,4]]

Code

java

import java.util.Arrays;
import java.util.Random;

class Solution {
    private final Random rand = new Random();

    public int[][] kClosest(int[][] points, int k) {
        int lo = 0, hi = points.length - 1;
        while (lo < hi) {
            int p = partition(points, lo, hi);
            if (p == k) {
                break; // indices [0, k) hold the k closest
            } else if (p < k) {
                lo = p + 1;
            } else {
                hi = p - 1;
            }
        }
        return Arrays.copyOfRange(points, 0, k);
    }

    private int partition(int[][] pts, int lo, int hi) {
        swap(pts, lo + rand.nextInt(hi - lo + 1), hi); // random pivot to the end
        int pivot = dist(pts[hi]);
        int store = lo;
        for (int i = lo; i < hi; i++) {
            if (dist(pts[i]) < pivot) {
                swap(pts, i, store++);
            }
        }
        swap(pts, store, hi); // pivot to its final position
        return store;
    }

    private int dist(int[] p) {
        return p[0] * p[0] + p[1] * p[1];
    }

    private void swap(int[][] pts, int i, int j) {
        int[] tmp = pts[i];
        pts[i] = pts[j];
        pts[j] = tmp;
    }
}

Complexity

Time: O(n) average — the range shrinks geometrically in expectation; O(n²) worst case with consistently bad pivots.
Space: O(1) extra (iterative), plus O(k) for the output.

Complexity Comparison

Approach	Time	Space
Sort by Distance	O(n log n)	O(log n)–O(n)
Max-Heap of Size k	O(n log k)	O(k)
Quickselect	O(n) average, O(n²) worst	O(1)

Common Pitfalls

Using Math.sqrt and doubles — unnecessary and invites precision bugs; compare x² + y² as ints (max 2 · 10^8, fits in int).
Using a min-heap capped at k — it would evict the closest points instead of the farthest.
Writing the comparator as dist(a) - dist(b) is fine here, but subtraction comparators overflow in general; Integer.compare is safer.
In Quickselect, using a fixed pivot (e.g. always hi) degrades to O(n²) on adversarial or sorted input — randomize.