K Closest Points to Origin

Master this topic with zero to advance depth.

K Closest Points to Origin

Given an array of points where points[i] = [xi, yi] 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 on the X-Y plane is the Euclidean distance $\sqrt{(x_1 - x_2)^2 + (y_1 - y_2)^2}$ .

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

Visual Representation

points = [[1,3],[-2,2]], k = 1

Distance of [1,3] = sqrt(1^2 + 3^2) = sqrt(10) ≈ 3.16
Distance of [-2,2] = sqrt((-2)^2 + 2^2) = sqrt(8) ≈ 2.83

Closest Point: [-2, 2]

Medium

Examples

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

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

Approach 1

Level I: Sorting (Naive)

Intuition

Calculate the squared distance for every point, then sort all points based on these distances. Return the first k points.

Thought Process

Use a custom comparator to compare points $(x_1, y_1)$ and $(x_2, y_2)$ based on $x_1^2 + y_1^2$ vs $x_2^2 + y_2^2$ .
Sort the array using this comparator.
Return the sub-array of size k.

⏱ O(N log N)💾 O(log N) to O(N) depending on sort implementation

Detailed Dry Run

points = [[1,3], [-2,2]], k = 1

Dist sq: [10, 8]
Sorted by dist: [[-2,2], [1,3]]
Take first k=1: [[-2,2]]

java

import java.util.Arrays;
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.copyOfRange(points, 0, k);
    }
}

Approach 2

Level II: Quickselect (Average Case Optimal)

Intuition

Similar to Kth Largest Element, we can use the Partition logic to find the $k$ closest points in $O(N)$ average time.

Thought Process

Define distance function $D(p) = x^2 + y^2$ .
Partition the points based on $D(p)$ .
If pivot index equals $k$ , all points to the left are the $k$ closest.
Recurse into the necessary partition.

⏱ O(N) Average, O(N^2) Worst Case💾 O(1) excluding recursion stack

Detailed Dry Run

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

Partition around pivot 20: [[3,3], [-2,4], [5,-1]]. Pivot 20 is at index 1.
k=2, so we need to ensure index 1 is part of the result. Return points[0...1].

java

import java.util.*;
class Solution {
    public int[][] kClosest(int[][] points, int k) {
        quickSelect(points, 0, points.length - 1, k);
        return Arrays.copyOfRange(points, 0, k);
    }
    void quickSelect(int[][] points, int L, int R, int k) {
        if (L >= R) return;
        int p = partition(points, L, R);
        if (p == k) return;
        if (p < k) quickSelect(points, p + 1, R, k);
        else quickSelect(points, L, p - 1, k);
    }
    int partition(int[][] points, int L, int R) {
        int[] pivot = points[R];
        int dist = dist(pivot), p = L;
        for (int i = L; i < R; i++) {
            if (dist(points[i]) <= dist) swap(points, i, p++);
        }
        swap(points, p, R);
        return p;
    }
    int dist(int[] p) { return p[0]*p[0] + p[1]*p[1]; }
    void swap(int[][] pts, int i, int j) { int[] t = pts[i]; pts[i] = pts[j]; pts[j] = t; }
}

Approach 3

Level III: Max-Heap (Optimal)

Intuition

Maintain a Max-Heap of size k to store the points with the smallest distances. If we find a point with a smaller distance than the maximum in our heap, we replace the maximum with it.

Thought Process

Initialize a Max-Heap based on squared Euclidean distance.
Iterate through all points.
Push the current point into the heap.
If heap size > k, pop the point with the largest distance (the root).
After all points are processed, the heap contains the k closest points.

⏱ O(N log K)💾 O(K)

Detailed Dry Run

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

Point [3,3], Dist 18: Heap [[3,3] (Dist 18)]
Point [5,-1], Dist 26: Heap [[5,-1] (Dist 26), [3,3] (Dist 18)]
Point [-2,4], Dist 20: Heap [[5,-1] (Dist 26), [-2,4], [3,3]] -> Pop [5,-1]. Final Heap: [[-2,4], [3,3]]

java

import java.util.*;

public class Main {
    public static int[][] kClosest(int[][] points, int k) {
        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.add(p);
            if (maxHeap.size() > k) maxHeap.poll();
        }
        int[][] res = new int[k][2];
        while (k > 0) res[--k] = maxHeap.poll();
        return res;
    }

    public static void main(String[] args) {
        int[][] res = kClosest(new int[][]{{3,3},{5,-1},{-2,4}}, 2);
        System.out.println(Arrays.deepToString(res));
    }
}

← Previous QuestionTask Scheduler

Next Question →Meeting Rooms II

Found an issue or have a suggestion?

Help us improve! Report bugs or suggest new features on our Telegram group.