Problem 18.2 · Heaps and Priority QueuesMedium

K Closest Points to Origin

PatternTop K with a Heap Divide and Conquer

What it teaches: K smallest by a computed key: a max-heap of size K keyed on squared distance.

Practise it on judges as “K Closest Points to Origin”.

The problem

Return the k points closest to (0, 0) by Euclidean distance, in any order.

Example 1

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

Squared distances 10 and 8.

Constraints

1 ≤ k ≤ n ≤ 10⁴
−10⁴ ≤ x, y ≤ 10⁴

Pattern clues in the wording

→ "K closest"
→ Any order

These clues point to Top K with a Heap: Keep a heap of size k: a min-heap for the k largest, a max-heap for the k smallest.

Stuck? Take one hint at a time

Solution.java · starter

import java.util.*;

class Solution {
    public int[][] kClosest(int[][] points, int k) {
        return new int[0][];
    }
}

Write your solution locally or in your editor for now. The in-browser runner (Java first, then Python, C++ and more) will run these tests right here.

Test cases

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

+ 1 hidden test the code runner will check

From slow to fast

Approaches

1. Max-heap of size k · O(n log k)2. Sort by distance · O(n log n)

Max-heap of size k

Time O(n log k) Space O(k)

For each point push it; if the heap exceeds k, poll the farthest.

Approach 1

import java.util.*;

class Solution {
    public int[][] kClosest(int[][] points, int k) {
        PriorityQueue<int[]> pq = new PriorityQueue<>((a, b) -> Integer.compare(d(b), d(a)));
        for (int[] p : points) {
            pq.offer(p);
            if (pq.size() > k) pq.poll();
        }
        return pq.toArray(new int[0][]);
    }

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

Verdict: Good for streams and small k.

Sort by distance

Time O(n log n) Space O(log n)

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

Approach 2

import java.util.*;

class Solution {
    public int[][] kClosest(int[][] points, int k) {
        Arrays.sort(points, Comparator.comparingInt(p -> p[0] * p[0] + p[1] * p[1]));
        return Arrays.copyOf(points, k);
    }
}

Verdict: Simplest; fine for one-off queries. Quickselect gives O(n) average.

Before you submit

Edge cases and common mistakes

Test these inputs

k = n
Ties in distance

Mistakes people make

Using a min-heap of all points (O(n) memory).
Computing Math.sqrt (slower, floating point).

Interview

Follow-up questions

Can you do better than O(n log k)?

Command Palette