We have a list of points on the plane.  Find the K closest points to the origin (0, 0).

(Here, the distance between two points on a plane is the Euclidean distance.)

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

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]].

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

Solution:

Create a distance array and then sort it.

Find the Kth value from the distance array.

Check the input array. If the distance is smaller or equal to the Kth value, add it into result.

class Solution {
    public int[][] kClosest(int[][] points, int K) {
        int n = points.length;
        int idx = 0;
        int[] dists = new int[n];
        int[][] res = new int[K][2];
        
        for (int i = 0; i < n; i++) {
            dists[i] = dist(points[i]);
        }
        
        Arrays.sort(dists);
        int distK = dists[K - 1];
        
        for (int i = 0; i < n; i++) {
            if (distK >= dist(points[i])) {
                res[idx++] = points[i];
            }
        }
        
        return res;
    }
    
    private int dist(int[] point) {
        return point[0] * point[0] + point[1] * point[1];
    }
}

Time complexity: O(n log(n))

Space complexity: O(n)