Leetcode: K Closest Points to Origin

This is an extremely common question at Facebook and Amazon. It is as follows: 

Given an array of points where points[i] = [xi, yi] represents a point on the X-Y plane and integer k, return the k closest points to the origin (0, 0). 

The distance between the 2 points is 

√(x1 - x2)^2 + (y1 - y2)^2).

Here's an example:

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

This is a very classical problem, the so-called Kth problem. 

The first solution is to sort the points by their distance directly, then get the top k closest points directly, but the solution is not very efficient, it's pretty slow. 

public int[][] kClosest(int[][] points, int K) { 

    Arrays.sort(points, (p1, p2) -> p1[0] * p1[0] + p1[1] * p1[1] - p2[0] * p2[0] - p2[1] * p2[1]);

    return Arrays.copyOfRange(points, 0, K);

}

Now there's a maximum heap solution for the Kth Closest Neighbors, where we sort by the closest values to the max heap. We make a maximum heap then iterate backward. Here's the code: 

class Solution{

public int[][] kClosest(int[][] points, int K) {

    PriorityQueue<int[]> pq = new PriorityQueue<int[]>((p1, p2) -> p2[0] * p2[0] + p2[1] * p2[1] - p1[0] * p1[0] - p1[1] * p1[1]);

    for (int[] p : points) {

        pq.offer(p);

        if (pq.size() > K) {

            pq.poll();

        }

    }

    int[][] res = new int[K][2];

    while (K > 0) {

        res[--K] = pq.poll();

    }

    return res;

}

}

Now the last solution I will discuss is the quick sort.

In quick sort, we choose a pivot and get all the elements smaller on the left side of the pivot and all the elements larger on the right side of the pivot. 

We then return the first K elements that are not greater than the pivot. Notice the other methods are private.

Here's the code: 

class Solution {

public int kClosest(int[][] points, int K) {

    int len = points.length;

    int l = 0;

    int r = len - 1;

    while(l <= r) {

        int mid = helper(points, l, r);

        if(mid == K) break;

        if(mid < K) {

            l = mid + 1;

        } else {

            r = mid - 1;

        }

    }

    return Arrays.copyOfRange(points, 0, K);

}

private int helper(int[][] A, int l, int r) {

    int[] pivot = A[l];

    while(l < r) {

        while(l < r && compare(A[r], pivot >= 0) r--;

        A[l] = A[r];

        while(l < r && compare(A[l], pivot) <= 0) l++;

        A[r] = A[l];

    }

    A[l] = pivot;

    return l;

}

private int compare(int[] p1, int[] p2) {

    return p1[0] * p1[0] + p1[1] * p1[1] - p2[0] * p2[0] - p2[1] * p2[1];

}

}