Problem

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 (i.e., √(x1 - x2)2 + (y1 - y2)2).

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

Examples

Example 1:

1
2
3
4
5
6
7

Input: points = [[1,3],[-2,2]], k = 1
Output: [[-2,2]]
Explanation:
The distance between (1, 3) and the origin is sqrt(10).
The distance between (-2, 2) and the origin is sqrt(8).
Since sqrt(8) < sqrt(10), (-2, 2) is closer to the origin.
We only want the closest k = 1 points from the origin, so the answer is just [[-2,2]].

Example 2:

1
2
3

Input: points = [[3,3],[5,-1],[-2,4]], k = 2
Output: [[3,3],[-2,4]]
Explanation: The answer [[-2,4],[3,3]] would also be accepted.

Notice the key requirement here: “K is much smaller than N. N is very large”. Definitely the brute-force solution by finding distance of all element and then sorting them in O(nlgn). But its costlier as we don’t need to sort all the n points as we are only concerned for first k points in the sorted list.

Follow up

Instead of origin, we are given some central point as well to the input? How will you handle that? -

Solution

Idea is to calculate the distance of the point from the origin and then use this problem - Kth Largest Element in an Array.

Method 1 - Using Max Heap

An efficient solution could be to use a max-heap of size k for maintaining k minimum distances.

A max heap has its largest element at the root. Whenever we add a point to the heap and its size exceeds K, we remove the root to discard the farthest point and retain the closest ones.

  • Use a max-heap to keep track of the k closest points.
  • Insert points into the heap based on their distance to the origin.
  • If the heap size exceeds k, remove the point with the maximum distance.
  • At the end, the heap contains the k closest points.

Code

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16

class Solution {
    public int[][] kClosest(int[][] points, int k) {
        PriorityQueue<int[]> heap = new PriorityQueue<>((a, b) -> (b[0] * b[0] + b[1] * b[1]) - (a[0] * a[0] + a[1] * a[1]));
        for (int[] point : points) {
            heap.offer(point);
            if (heap.size() > k) {
                heap.poll();
            }
        }
        int[][] ans = new int[k][2];
        for (int i = 0; i < k; i++) {
            ans[i] = heap.poll();
        }
        return ans;
    }
}
 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41

class Solution {

    static class Point {
        public int x;
        public int y;

        public Point(final int x, final int y) {
            this.x = x;
            this.y = y;
        }

        // distance from origin
        public long getDist() {
            return (long) x * x + y * y;
        }
    }

    public int[][] kClosest(int[][] points, final int k) {
        // Max heap
        final PriorityQueue<Point> maxHeap = new PriorityQueue<>((o1, o2) -> Long.compare(o2.getDist(), o1.getDist()));

        for (int[] point : points) {
            Point p = new Point(point[0], point[1]);
            maxHeap.offer(p);
            if (maxHeap.size() > k) {
                maxHeap.poll();
            }
        }

        int[][] ans = new int[k][2];
        int index = 0;
        while (!maxHeap.isEmpty()) {
            Point p = maxHeap.poll();
            ans[index][0] = p.x;
            ans[index][1] = p.y;
            index++;
        }

        return ans;
    }
}
 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12

class Solution:
    def kClosest(self, points: List[List[int]], k: int) -> List[List[int]]:
        max_heap = []
        for (x, y) in points:
            dist = -(x**2 + y**2)  # Use negative for max-heap
            if len(max_heap) < k:
                heapq.heappush(max_heap, (dist, x, y))
            else:
                heapq.heappushpop(max_heap, (dist, x, y))
        
        ans = [[x, y] for (dist, x, y) in max_heap]
        return ans

Complexity

  • ⏰ Time complexity: O(n log k) where n is the number of points and k is the number of closest points to find.
  • 🧺 Space complexity: O(k) to store the k closest points.

Method 2 -Quickselect

Using Quickselect, we can select the kth minimum distance from the list of distances in O(n) time. Then, we need to go through the array of points and output the first k elements with a distance less than or equal to the kth minimum distance. This is an O(n) time in-place algorithm with constant space.

Approach

  1. Distance Calculation: For each point, calculate the squared Euclidean distance to the origin. The squared distance avoids dealing with floating-point precision issues.
  2. Quickselect: Use the Quickselect algorithm to partition the points such that the k closest points are in the correct positions.
  3. Partition: After applying Quickselect, the first k points in the array are the closest points.
  4. Output: Return these k closest points.

Code

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47

class Solution {
    public int[][] kClosest(int[][] points, int k) {
        quickSelect(points, 0, points.length - 1, k);
        int[][] ans = new int[k][2];
        System.arraycopy(points, 0, ans, 0, k);
        return ans;
    }

    private void quickSelect(int[][] points, int left, int right, int k) {
        if (left < right) {
            int pivotIndex = partition(points, left, right);
            if (pivotIndex == k - 1) {
                return;
            } else if (pivotIndex < k - 1) {
                quickSelect(points, pivotIndex + 1, right, k);
            } else {
                quickSelect(points, left, pivotIndex - 1, k);
            }
        }
    }

    private int partition(int[][] points, int left, int right) {
        int[] pivot = points[right];
        int pivotDist = distance(pivot);
        int swapIndex = left;

        for (int i = left; i < right; i++) {
            if (distance(points[i]) <= pivotDist) {
                swap(points, i, swapIndex);
                swapIndex++;
            }
        }

        swap(points, swapIndex, right);
        return swapIndex;
    }

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

    private int distance(int[] point) {
        return point[0] * point[0] + point[1] * point[1];
    }
}
 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28

class Solution:
    def kClosest(self, points: List[List[int]], k: int) -> List[List[int]]:
        def dist(point: List[int]) -> int:
            return point[0] ** 2 + point[1] ** 2

        def partition(left: int, right: int) -> int:
            pivot = points[right]
            pivot_dist = dist(pivot)
            i = left
            for j in range(left, right):
                if dist(points[j]) <= pivot_dist:
                    points[i], points[j] = points[j], points[i]
                    i += 1
            points[i], points[right] = points[right], points[i]
            return i

        def quickSelect(left: int, right: int, k: int) -> None:
            if left < right:
                pivot_index = partition(left, right)
                if pivot_index == k - 1:
                    return
                elif pivot_index < k - 1:
                    quickSelect(pivot_index + 1, right, k)
                else:
                    quickSelect(left, pivot_index - 1, k)

        quickSelect(0, len(points) - 1, k)
        return points[:k]

Complexity

  • ⏰ Time complexity: O(n) where n is the number of points (average case for Quickselect)
  • 🧺 Space complexity: O(1) for the in-place partitioning.

Follow up - K closest distance from central point

We just need to update the distance calculation -

1
2
3
4

 private long distance(int[] point, int[] centralPoint) {
	 return (long) (point[0] - centralPoint[0]) * (point[0] - centralPoint[0]) +
			 (point[1] - centralPoint[1]) * (point[1] - centralPoint[1]);
 }

Code

1
2


##### Java
 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28

class Solution {
    public int[][] kClosest(int[][] points, int[] centralPoint, int k) {
        // Max heap to keep the k closest points
        PriorityQueue<int[]> maxHeap = new PriorityQueue<>((o1, o2) -> Long.compare(
                distance(o2, centralPoint), distance(o1, centralPoint)
        ));

        for (int[] point : points) {
            maxHeap.offer(point);
            if (maxHeap.size() > k) {
                maxHeap.poll();
            }
        }

        int[][] ans = new int[k][2];
        int index = 0;
        while (!maxHeap.isEmpty()) {
            ans[index++] = maxHeap.poll();
        }

        return ans;
    }

    private long distance(int[] point, int[] centralPoint) {
        return (long) (point[0] - centralPoint[0]) * (point[0] - centralPoint[0]) +
                (point[1] - centralPoint[1]) * (point[1] - centralPoint[1]);
    }
}
 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14

class Solution:
    def kClosest(self, points: List[List[int]], central_point: List[int], k: int) -> List[List[int]]:
        def dist(point: List[int]) -> int:
            return (point[0] - central_point[0]) ** 2 + (point[1] - central_point[1]) ** 2

        max_heap = []
        for point in points:
            distance = dist(point)
            if len(max_heap) < k:
                heapq.heappush(max_heap, (-distance, point))
            else:
                heapq.heappushpop(max_heap, (-distance, point))
        
        return [point for (_, point) in max_heap]