K Closest Points to Origin

1. Problem & Core Objective

Given an array of points on a plane and an integer k, return the k points closest to the origin (0, 0). The answer may be returned in any order.

points = [[1,3],[-2,2]], k = 1     →  [[-2,2]]
         distances √10 and √8

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

Constraints: 1 <= k <= points.length <= 10^4 · -10^4 <= x, y <= 10^4

What's actually being tested: two things. That you skip the square root — comparing distances doesn't need it — and that you pick the right selection strategy. It's question 1's heap inversion applied to the k smallest, so the heap flips the other way.

2. First-Principles Thought Process

Don't compute the square root

Distance from the origin is √(x² + y²). But we only ever compare distances, never report them.

√ is monotonically increasing, so for non-negative a and b:

√a < √b   if and only if   a < b

So comparing x² + y² gives identical ordering. Dropping the root removes a floating-point operation per comparison and, more importantly, keeps everything in integers — no precision concerns, no ties that are equal in reality but differ in the last bit.

With |x|, |y| <= 10^4, the maximum squared distance is 2 × 10^8, comfortably inside int. Worth checking rather than assuming — this is exactly the bound that decides whether you need long.

Which heap, and why it flips

We want the k smallest distances. By the inversion from question 1:

• keep the k smallest in a heap of size k
• the one to evict when a closer point arrives is the largest of those kept
• so the root must be the maximum → max-heap

Min-heap for the k largest; max-heap for the k smallest. This is the mirror of question 1, and stating the rule both ways is the cleanest way to show it's understood rather than memorised.

The alternatives are genuinely competitive

Unlike question 1, this isn't a stream — the whole array is in memory. That opens two more options:

The heap wins when k is much smaller than n; quickselect wins on raw speed for one query. Being able to say which and why is the answer, not picking one.

3. Solution Paths

Approach 1 — Sort everything by distance (brute force, and perfectly reasonable)

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

• Time O(n log n) · Space O(1) extra (sorts in place) — O(log n) for the sort's stack

Counter-questions on this approach

⭐ "Two lines and it's correct. What's the objection?"

It computes the full ordering of all n points to read the first k. If k = 1 I've sorted 10,000 points to find one.

With a size-k heap it's O(n log k); with quickselect, O(n) average. Neither produces the wasted ordering.

That said, at n = 10^4 this runs in about 1.3 × 10^5 comparisons — instant. I'd write it first and say it's what I'd ship unless k ≪ n or the data streams.

⭐ "Is p[0] * p[0] + p[1] * p[1] safe from overflow?"

Yes, and it's worth confirming rather than assuming. With coordinates bounded by 10^4, each square is at most 10^8 and the sum at most 2 × 10^8 — int holds up to about 2.1 × 10^9. Safe with an order of magnitude to spare.

If coordinates could reach 10^5, the sum would be 2 × 10^10 and I'd need long. That's the check to do before writing the comparator.

"Why comparingInt rather than (a, b) -> dist(a) - dist(b)?"

Subtraction overflows when the difference exceeds int range. Here distances are bounded so it happens to be safe, but comparingInt is both safer and clearer. It also computes the key once per comparison rather than twice, though the JIT usually handles that.

"Does this mutate the input?"

Yes — Arrays.sort sorts in place, so the caller's array is reordered. Worth flagging; if that matters, sort a copy of the indices instead.

Approach 2 — Max-heap of size k (optimal when k ≪ n)

public int[][] kClosest(int[][] points, int k) {
    // max-heap on squared distance: the root is the FARTHEST of the k kept
    PriorityQueue<int[]> maxHeap = new PriorityQueue<>(
        (a, b) -> (b[0]*b[0] + b[1]*b[1]) - (a[0]*a[0] + a[1]*a[1]));

    for (int[] p : points) {
        maxHeap.offer(p);
        if (maxHeap.size() > k) maxHeap.poll();     // drop the farthest
    }

    int[][] result = new int[k][];
    for (int i = 0; i < k; i++) result[i] = maxHeap.poll();
    return result;
}

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

Result: [-2,4] and [3,3] ✓

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

Counter-questions on this approach

⭐ "Why a max-heap when the question asks for the closest?"

Because I'm keeping the k smallest distances, and the element to evict when a closer point arrives is the largest of those kept. A max-heap puts exactly that at the root.

It's the mirror of the previous stream question: there I kept the k largest and needed a min-heap. The rule is min-heap for the k largest, max-heap for the k smallest — the heap is always inverted relative to the wording, because the root is the eviction candidate, not the answer.

⭐ "This is O(n log k) versus sorting's O(n log n). When does that actually matter?"

When k ≪ n. At n = 10^4 and k = 10, log k is about 3 versus log n about 13 — roughly a 4× improvement, which is real but not dramatic at this size.

Where it matters much more is space and streaming. The heap holds k points, not n, so it works on a stream you can't store and on datasets that don't fit in memory. Sorting requires the whole array at once.

"Your comparator uses subtraction. You just warned against that."

Fair — and here the operands are squared distances up to 2 × 10^8, so a difference could be 2 × 10^8, which is safe but uncomfortably close to needing thought. The cleaner form is:

Java

Comparator.comparingInt((int[] p) -> p[0]*p[0] + p[1]*p[1]).reversed()

That's overflow-proof by construction and states "farthest first" directly. I'd write that version.

"You recompute the squared distance on every comparison. Is that wasteful?"

Slightly — each comparison does four multiplications. For n = 10^4 with log k comparisons each, it's negligible. If it mattered you'd store {dist, x, y} triples so the distance is computed once per point, trading a little memory for arithmetic.

"Does the output order matter?"

The problem says any order, which is why draining the heap is fine even though it yields farthest-to-closest. If sorted output were required I'd reverse the drained array, or use a min-heap over all n points and poll k times.

Approach 3 — Quickselect

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 - 1) break;
        if (p < k - 1) lo = p + 1; else hi = p - 1;
    }
    return Arrays.copyOfRange(points, 0, k);
}

private int partition(int[][] pts, int lo, int hi) {
    int pivot = dist(pts[hi]), store = lo;
    for (int i = lo; i < hi; i++)
        if (dist(pts[i]) < pivot) swap(pts, i, store++);
    swap(pts, store, hi);
    return store;
}

private int dist(int[] p) { return p[0]*p[0] + p[1]*p[1]; }
private void swap(int[][] a, int i, int j) { int[] t = a[i]; a[i] = a[j]; a[j] = t; }

• Time O(n) average, O(n²) worst · Space O(1)

Counter-questions on this approach

⭐ "Why is this O(n) when sorting is O(n log n)?"

Because it only recurses into the side that contains the answer. Sorting orders both halves; quickselect discards one entirely at each step.

The work is n + n/2 + n/4 + … = 2n on average, so O(n). The trade is that the array ends up partially ordered — the first k are the closest, but in no particular order among themselves — which is exactly what the problem asks for and nothing more.

⭐ "What's the worst case, and does it matter?"

O(n²), when every pivot is the minimum or maximum — for instance, an already-sorted array with last-element pivots. An adversary who knows your pivot rule can force it.

The standard defence is a random pivot, which makes the bad case astronomically unlikely, or median-of-medians, which guarantees O(n) at a much worse constant. In an interview I'd note the risk and use a random pivot.

"Does it mutate the input?"

Heavily — it reorders the array in place. That's how it achieves O(1) space, and it's a real cost to mention.

"Which would you actually submit?"

The heap or the sort, and I'd say why: quickselect is the fastest on average but is the most code, has a bad worst case, and destroys the input ordering. At n = 10^4 the sort runs instantly. I'd write the heap if asked about streaming or k ≪ n, and mention quickselect as the O(n) option.

Comparison

4. Why the Optimal Wins

There isn't a single winner here, and saying so is better than pretending otherwise.

Against sorting: both the heap and quickselect avoid computing the full order when only a prefix is wanted. Sorting to read k of n items is doing n log n work for a k-sized answer.

Heap versus quickselect: the heap is O(n log k) with O(k) space and works on a stream; quickselect is O(n) average with O(1) space but needs the whole array and reorders it. For one query on an in-memory array, quickselect. For a stream or k ≪ n, the heap.

The framing worth keeping:

Comparing distances never needs the square root — x² + y² orders identically and stays in integers. And for the k smallest, the heap is a MAX-heap, because its root is the eviction candidate.

5. Java Prerequisites

Heap of arrays with a safe comparator

PriorityQueue<int[]> maxHeap = new PriorityQueue<>(
    Comparator.comparingInt((int[] p) -> p[0]*p[0] + p[1]*p[1]).reversed());

The explicit (int[] p) type is needed — without it, type inference fails on the chained reversed().

reversed() reverses the whole chain, not just the last key. With a single key that's fine; with thenComparing it's a classic trap. See 04.

Overflow check before writing the comparator. |coord| <= 10^4 → squared sum <= 2 × 10^8 → int is safe. At 10^5 it would not be.

Arrays.copyOfRange(arr, 0, k) returns the first k elements as a new array.

Arrays.sort on objects is a stable merge sort, O(n log n); on primitives it's dual-pivot quicksort and not stable.

6. Interview Communication Guide

Clarifying questions: Does output order matter (no — it permits the heap drain)? Can points repeat (yes; duplicates are fine)? Is the input array mutable (sorting and quickselect both reorder it)? Is this one query or a stream (it decides heap vs quickselect)? What's the coordinate range (10^4 — it decides int vs long)?

The pitch

"First, I'd skip the square root. I only ever compare distances, never report them, and √ is monotonically increasing — so x² + y² gives the identical ordering. That keeps everything in integers, avoiding floating point entirely.

Worth checking the bound: coordinates up to 10^4 give a squared sum up to 2 × 10^8, which fits in int with room to spare. At 10^5 I'd need long.

Then there are three reasonable approaches, and which is best depends on the setting.

Sorting by distance and taking the first k is two lines and O(n log n). At n = 10^4 that's instant, and it's what I'd ship absent other requirements.

A max-heap of size k is O(n log k) and O(k) space. Max-heap, not min — I'm keeping the k smallest distances, so the element to evict when a closer point arrives is the largest of those kept, and that's what I need at the root. It's the mirror of the streaming question: min-heap for the k largest, max-heap for the k smallest. The heap is always inverted, because the root is the eviction candidate rather than the answer.

Quickselect is O(n) on average — it partitions and recurses only into the side containing the answer, so it never orders the rest. The worst case is O(n²) on adversarial pivots, which a random pivot defends against.

So: heap if the data streams or k is much smaller than n; quickselect for one query on an in-memory array; sorting if clarity matters more than the constant."

Edge cases to volunteer:

Name the equidistant case and the coordinate extreme. The first confirms you read "any order"; the second is where an implementation using long unnecessarily — or int incorrectly, at larger bounds — shows up.

7. Follow-Up Questions — Modified Constraints

⭐ "The points arrive as a stream too large to store."

Quickselect and sorting are both out — they need the whole array. The max-heap is the only option: it holds k points regardless of stream length, at O(log k) per point. This is precisely where the heap's O(k) space stops being a minor detail.

⭐ "Find the k closest points to an arbitrary point (a, b), not the origin."

Replace x² + y² with (x−a)² + (y−b)². Nothing structural changes — but re-check overflow: with coordinates and the target both at 10^4, a difference can be 2 × 10^4 and its square 4 × 10^8, so the sum reaches 8 × 10^8. Still int, but the margin has shrunk fourfold. That's the kind of thing worth recomputing rather than carrying over.

"Answer many k-closest queries against different targets."

Build a k-d tree or ball tree once, O(n log n), then each query is O(log n + k) on average. That's the right structure when queries dominate — a heap would rescan all n points per query.

"Return them sorted by distance."

Drain the max-heap and reverse, O(k log k). Or use a min-heap over all n and poll k times, O(n + k log n) with O(n) heapify — which is better when k is close to n.

"What if the points were in 3D, or d dimensions?"

The comparator becomes a sum over d coordinates. Everything else is unchanged, but overflow gets closer: d terms each up to 10^8 means long becomes necessary around d = 20. In high dimensions the k-d tree also degrades toward linear scan — the curse of dimensionality — so approximate methods like LSH take over.

"Find the k closest pairs among the points, not to a fixed origin."

Different problem, and much harder. Naively O(n²) pairs fed through a size-k heap. Closest-pair-of-points has a classic O(n log n) divide-and-conquer solution; the k-closest-pairs generalisation is harder still. Worth naming that it isn't a variation of this technique.