K Closest Points to Origin

Problem Description

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

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 <= 104
• -104 <= xi, yi <= 104

Solution (JavaScript)

/**
 * @param {number[][]} points
 * @param {number} k
 * @return {number[][]}
 */
var kClosest = function(points, k) {
    let len = points.length, l = 0, r = len - 1;
    while (l <= r) {
        let mid = helper(points, l , r);
        if (mid === k) break;
        if (mid < k) {
            l = mid + 1;
        } else {
            r = mid - 1;
        }
    }
    return points.slice(0, k);
};

function helper(A, l , r) {
    let pivot = A[l];
    while (l < r) {
        while(l < r && closerToOrigin(pivot, A[r])) r--;
        A[l] = A[r];
        while(l < r && closerToOrigin(A[l], pivot)) l++;
        A[r] = A[l];
    }
    A[l] = pivot;
    return l;
}

function closerToOrigin(p1, p2) {
    return p1[0]**2 + p1[1]**2 - p2[0]**2 - p2[1]**2 <= 0;
}