Definition
- Max-heap of size K – a priority queue that keeps the farthest point on top; when a closer point arrives and the heap exceeds K, the farthest is evicted.
- Squared distance – comparing x*x + y*y avoids the square root; squaring preserves ordering for non-negative values.
- Time O(n log K) – each of n points is offered to a heap of size at most K, so each offer is O(log K).
- Space O(K) – the heap never holds more than K entries.
- QuickSelect alternative – average O(n) time using the selection algorithm on squared distances.
The National Bank of India mobile app shows the K nearest ATMs to the customer's location. You keep a shortlist of K, ordered so the farthest sits on top. When you consider a new ATM, you add it and, if the list grew past K, drop the current farthest. The K that remain are the nearest.
Understanding the Concepts
What is the K Closest Points problem?
Given n points on a 2D plane and an integer K, return the K points closest to the origin (0,0). Distance is Euclidean: sqrt(x^2 + y^2). The output order does not matter.
At the National Bank of India, the mobile app finds the K nearest ATMs or branches to the customer's GPS coordinates.
package com.bank.geo;
import java.util.*;
public class KClosestPointsDemo {
public static void main(String[] args) {
int[][] points = {{3, 3}, {5, -1}, {-2, 4}, {1, 1}, {-1, 0}};
int k = 3;
int[][] result = kClosest(points, k);
System.out.println("K closest: " + Arrays.deepToString(result));
}
static int[][] kClosest(int[][] points, int k) {
// Max-heap by squared distance (farthest on top)
PriorityQueue<int[]> maxHeap = new PriorityQueue<>(
(a, b) -> Long.compare(dist(b), dist(a))
);
for (int[] p : points) {
maxHeap.offer(p);
if (maxHeap.size() > k) {
maxHeap.poll(); // remove farthest
}
}
int[][] result = new int[k][2];
int i = 0;
while (!maxHeap.isEmpty()) {
result[i++] = maxHeap.poll();
}
return result;
}
static long dist(int[] p) {
return (long) p[0] * p[0] + (long) p[1] * p[1];
}
}
K closest: [[1, 1], [-1, 0], [-2, 4]]
Why use squared distance instead of true Euclidean distance?
Squaring is monotonic for non-negative values: if sqrt(a) < sqrt(b), then a < b. So comparing squared distances gives the same ordering as comparing true distances, without the cost and floating-point imprecision of square roots. Use long for the squared sum to avoid overflow on large coordinates.
// WRONG: floating-point square root (imprecise, slow)
double d1 = Math.sqrt(x1*x1 + y1*y1);
double d2 = Math.sqrt(x2*x2 + y2*y2);
// CORRECT: squared distance (exact, fast)
long d1 = (long) x1 * x1 + (long) y1 * y1;
long d2 = (long) x2 * x2 + (long) y2 * y2;
// d1 < d2 has the same result as sqrt(d1) < sqrt(d2)
Why a max-heap of size K instead of a min-heap of all n?
A min-heap of all n points requires O(n) space and O(n log n) to extract K elements. A max-heap capped at K uses only O(K) space and O(n log K) total time. When n is large and K is small, the difference is significant.
| Approach | Time | Space |
|----------|------|-------|
| Sort all n points | O(n log n) | O(n) |
| Min-heap of all n | O(n + K log n) | O(n) |
| Max-heap of size K | O(n log K) | O(K) |
| QuickSelect | O(n) average | O(n) |
// Max-heap advantage: for n=1,000,000 and K=10
// Max-heap: 1,000,000 * log(10) = ~3.3M operations
// Full sort: 1,000,000 * log(1,000,000) = ~20M operations
How does the max-heap eviction work step by step?
The max-heap keeps the farthest point on top. When a new point is offered and the heap exceeds K, the farthest (top) is polled. After processing all points, the heap contains exactly the K closest.
package com.bank.geo;
import java.util.*;
public class HeapStepByStep {
public static void main(String[] args) {
int[][] points = {{3, 3}, {5, -1}, {-2, 4}};
int k = 2;
PriorityQueue<int[]> maxHeap = new PriorityQueue<>(
(a, b) -> Long.compare(dist(b), dist(a)));
for (int[] p : points) {
maxHeap.offer(p);
System.out.println("Added " + Arrays.toString(p)
+ " (dist=" + dist(p) + "), heap size=" + maxHeap.size());
if (maxHeap.size() > k) {
int[] evicted = maxHeap.poll();
System.out.println(" Evicted " + Arrays.toString(evicted)
+ " (dist=" + dist(evicted) + ")");
}
}
System.out.println("Result: ");
while (!maxHeap.isEmpty()) {
System.out.println(" " + Arrays.toString(maxHeap.poll()));
}
}
static long dist(int[] p) {
return (long) p[0] * p[0] + (long) p[1] * p[1];
}
}
Added [3, 3] (dist=18), heap size=1
Added [5, -1] (dist=26), heap size=2
Added [-2, 4] (dist=20), heap size=3
Evicted [5, -1] (dist=26)
Result:
[-2, 4]
[3, 3]
What is the QuickSelect alternative?
QuickSelect partitions the array around a pivot's squared distance. Elements closer than the pivot end up on one side, and farther on the other. By recursing into the correct partition, you find the K-th closest in average O(n) time. However, the worst case is O(n^2).
package com.bank.geo;
import java.util.Arrays;
public class QuickSelectClosest {
static int[][] kClosest(int[][] points, int k) {
quickSelect(points, 0, points.length - 1, k);
return Arrays.copyOf(points, k);
}
static void quickSelect(int[][] pts, int lo, int hi, int k) {
if (lo >= hi) return;
int pivot = partition(pts, lo, hi);
if (pivot == k) return;
if (pivot < k) quickSelect(pts, pivot + 1, hi, k);
else quickSelect(pts, lo, pivot - 1, k);
}
static int partition(int[][] pts, int lo, int hi) {
long pivotDist = dist(pts[hi]);
int i = lo;
for (int j = lo; j < hi; j++) {
if (dist(pts[j]) <= pivotDist) {
int[] temp = pts[i]; pts[i] = pts[j]; pts[j] = temp;
i++;
}
}
int[] temp = pts[i]; pts[i] = pts[hi]; pts[hi] = temp;
return i;
}
static long dist(int[] p) {
return (long) p[0] * p[0] + (long) p[1] * p[1];
}
public static void main(String[] args) {
int[][] points = {{3, 3}, {5, -1}, {-2, 4}, {1, 1}, {-1, 0}};
int[][] result = kClosest(points, 3);
System.out.println("K closest: " + Arrays.deepToString(result));
}
}
K closest: [[-1, 0], [1, 1], [-2, 4]]
What is the complete banking implementation?
Here is the ATM finder using the max-heap approach with proper banking context.
package com.bank.geo;
import java.util.*;
public class NearestATMFinder {
record ATM(String id, int x, int y) {}
public List<ATM> findNearest(List<ATM> atms, int customerX, int customerY, int k) {
PriorityQueue<ATM> maxHeap = new PriorityQueue<>((a, b) -> {
long da = squaredDist(a.x - customerX, a.y - customerY);
long db = squaredDist(b.x - customerX, b.y - customerY);
return Long.compare(db, da); // farthest on top
});
for (ATM atm : atms) {
maxHeap.offer(atm);
if (maxHeap.size() > k) {
maxHeap.poll();
}
}
List<ATM> result = new ArrayList<>(maxHeap);
return result;
}
private long squaredDist(int dx, int dy) {
return (long) dx * dx + (long) dy * dy;
}
public static void main(String[] args) {
NearestATMFinder finder = new NearestATMFinder();
List<ATM> atms = List.of(
new ATM("ATM-001", 1, 2),
new ATM("ATM-002", 5, 8),
new ATM("ATM-003", -1, 0),
new ATM("ATM-004", 3, 3),
new ATM("ATM-005", 10, 10)
);
List<ATM> nearest = finder.findNearest(atms, 0, 0, 3);
nearest.forEach(a -> System.out.println(a.id() + " at (" + a.x() + "," + a.y() + ")"));
}
}
ATM-003 at (-1,0)
ATM-001 at (1,2)
ATM-004 at (3,3)
What are common mistakes?
The most common mistake is using a min-heap instead of a max-heap, which requires inserting all n points and extracting K, using O(n) space instead of O(K).
// WRONG: min-heap of all n points (O(n) space)
PriorityQueue<int[]> minHeap = new PriorityQueue<>(
(a, b) -> Long.compare(dist(a), dist(b))); // closest on top
for (int[] p : points) minHeap.offer(p); // all n in heap
int[][] result = new int[k][2];
for (int i = 0; i < k; i++) result[i] = minHeap.poll(); // extract K
// CORRECT: max-heap of size K (O(K) space)
PriorityQueue<int[]> maxHeap = new PriorityQueue<>(
(a, b) -> Long.compare(dist(b), dist(a))); // farthest on top
for (int[] p : points) {
maxHeap.offer(p);
if (maxHeap.size() > k) maxHeap.poll(); // evict farthest
}
What should a senior engineer know?
- Generalization. This is a K-th smallest problem. Any "top K" or "K-th element" problem can use the same max-heap-of-size-K pattern.
- Custom origin. To find K closest to a customer's location instead of the origin, subtract the customer's coordinates before computing squared distance.
- Tie-breaking. If two points have the same distance, any tie-breaking is acceptable unless the problem specifies otherwise.
- Streaming data. The max-heap approach works for streaming data: you do not need all points in memory at once, just the current K best.
Visual Overview
Points: (3,3) (5,-1) (-2,4) (1,1) (-1,0) K=3
Squared distances from origin:
(3,3) -> 18
(5,-1) -> 26
(-2,4) -> 20
(1,1) -> 2
(-1,0) -> 1
Max-heap (farthest on top), cap at K=3:
Step 1: offer (3,3)=18 heap: [18]
Step 2: offer (5,-1)=26 heap: [26, 18]
Step 3: offer (-2,4)=20 heap: [26, 18, 20] size=3=K
Step 4: offer (1,1)=2 heap: [26, 18, 20, 2] size=4>K
poll top=26 (5,-1) heap: [20, 18, 2]
Step 5: offer (-1,0)=1 heap: [20, 18, 2, 1] size=4>K
poll top=20 (-2,4) heap: [18, 2, 1]
Result: (3,3)=18, (1,1)=2, (-1,0)=1 -> K closest
Key Takeaway
K Closest Points to Origin asks for the K nearest points and the efficient solution uses a max-heap of size K ordered by squared distance, with the farthest point on top. I offer each point, and when the heap exceeds K I poll the top to remove the current farthest. After all points are processed, the heap holds exactly the K closest.
Comparing squared distances instead of true Euclidean distances is a deliberate choice. Squaring is monotonic for non-negative values, so it preserves ordering, and it avoids both the cost and the floating-point imprecision of square roots. I use long for the squared sum to prevent overflow on large coordinates.
The choice between a max-heap of size K and other approaches is what interviewers probe. Capping the heap at K bounds both memory and per-element cost at O(log K), which matters when n is large and K is small. QuickSelect gives average O(n) time but has O(n^2) worst case and requires all points in memory.
At the National Bank of India, the mobile app used this pattern to show the K nearest ATMs to the customer's GPS coordinates. The service scanned the ATM network once, held only K candidates in memory, and returned the nearest ATMs quickly without sorting the entire network or computing square roots.
Remember one sentence: a max-heap of size K keyed on squared distance finds the K closest points in O(n log K) time and O(K) space by evicting the farthest point whenever the heap overflows.
---