[leetcode #973] K Closest Points to Origin

Problem

Given an array of points where points[i] = [xᵢ, yᵢ] 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)² + (y1 - y2)²).

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

Example 1:

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:

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.

Constraints:

・ 1 <= k <= points.length <= 10⁴
・ -10⁴ < xᵢ, yᵢ < 10⁴

Idea

간단하게 생각할 수 있는 방법은 원점으로부터 거리를 구한 뒤, 해당 거리로 정렬을 하는 것이다.

정렬한 뒤 index를 사용해 주어진 points를 array로 만들어 리턴하면 된다. (O(NlogN))

그런데 Solution을 보면 O(N)으로 풀 수 있는 방법이 있다. Binary Search나 Quick Select를 이용하는 방법인데, 시간이 날 때 읽어봐야겠다.

Solution

class Solution {
    public int[][] kClosest(int[][] points, int k) {
        int[][] squares = new int[points.length][2];

        for (int i=0; i < points.length; i++) {
            squares[i][0] = i;
            squares[i][1] = points[i][0] * points[i][0] + points[i][1] * points[i][1];
        }

        Arrays.sort(squares, (int[] a, int[] b) -> { return a[1] - b[1];});

        List<int[]> res = new ArrayList<>();
        for (int i=0; i < k; i++) {
            int index = squares[i][0];
            res.add(new int[]{points[index][0], points[index][1]});
        }

        return res.toArray(new int[res.size()][2]);
    }
}

Quick Select는 다음과 같다.

class Solution {
    public int[][] kClosest(int[][] points, int k) {
        // Precompute the Euclidean distance for each point,
        // define the initial binary search range,
        // and create a reference list of point indices
        double[] distances = new double[points.length];
        double low = 0, high = 0;
        List<Integer> remaining = new ArrayList<>();
        for (int i = 0; i < points.length; i++) {
            distances[i] = euclideanDistance(points[i]);
            high = Math.max(high, distances[i]);
            remaining.add(i);
        }

        // Perform a binary search of the distances
        // to find the k closest points
        List<Integer> closest = new ArrayList<>();
        while (k > 0) {
            double mid = low + (high - low) / 2;
            List<List<Integer>> result = splitDistances(remaining, distances, mid);
            List<Integer> closer = result.get(0), farther = result.get(1);
            if (closer.size() > k) {
                // If more than k points are in the closer distances
                // then discard the farther points and continue
                remaining = closer;
                high = mid;
            } else {
                // Add the closer points to the answer array and keep
                // searching the farther distances for the remaining points
                k -= closer.size();
                closest.addAll(closer);
                remaining = farther;
                low = mid;
            }
        }

        // Return the k closest points using the reference indices
        k = closest.size();
        int[][] answer = new int[k][2];
        for (int i = 0; i < k; i++) {
            answer[i] = points[closest.get(i)];
        }
        return answer;
    };

    private List<List<Integer>> splitDistances(List<Integer> remaining, double[] distances,
                                               double mid) {
        // Split the distances around the midpoint
        // and return them in separate lists
        List<List<Integer>> result = new ArrayList<>();
        List<Integer> closer = new ArrayList<>();
        List<Integer> farther = new ArrayList<>();
        result.add(closer);
        result.add(farther);
        for (int point : remaining) {
            if (distances[point] <= mid) {
                closer.add(point);
            } else {
                farther.add(point);
            }
        }
        return result;
    }

    private double euclideanDistance(int[] point) {
        // Calculate and return the Euclidean distance
        return point[0] * point[0] + point[1] * point[1];
    }
}

Reference

https://leetcode.com/problems/k-closest-points-to-origin/
https://leetcode.com/problems/k-closest-points-to-origin/solution/