973. K Closest Points to Origin

K Closest Points to Origin - LeetCode

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


  • code
class Solution:
    def kClosest(self, points: List[List[int]], k: int) -> List[List[int]]:
        return heapq.nsmallest(k, points, key = lambda x: x[0] ** 2 + x[1] ** 2)
  • code max #priorityQueue
class Solution {
    
    public int distance(int[] point){
        return point[0] * point[0] + point[1] * point[1];
    }
    
    public int[][] kClosest(int[][] points, int k) {
        PriorityQueue<int[]> pq = new PriorityQueue<>((a, b) -> b[1] - a[1]);
        for (int i = 0; i < k; i++){
            pq.add(new int[]{i, distance(points[i])});
        }
        for (int i = k; i < points.length; i++){
            if (distance(points[i]) < pq.peek()[1]){
                pq.poll();
                pq.add(new int[]{i, distance(points[i])});
            }
        }
        int[][] ans = new int[k][2];
        for (int i = 0; i < k; i++){
            ans[i] = points[pq.poll()[0]];
        }
        return ans;
    }
}
  • code
class Solution:
    def kClosest(self, points: List[List[int]], k: int) -> List[List[int]]:
        def dis_negative(x, y): # negate to make a max heap, python default min heap
            return -(x ** 2 + y ** 2)
        heap = [(dis_negative(*points[i]), i) for i in range(k)]
        heapq.heapify(heap)
        for i in range(k, len(points)):
            dist = dis_negative(*points[i])
            if dist > heap[0][0]:
                heapq.heappushpop(heap, (dist, i))
        
        return [points[i] for distance, i in heap]
  • code #binarySearch the mid distance to split the input, store index
class Solution:
    def kClosest(self, points: List[List[int]], k: int) -> List[List[int]]:
        def dist(x, y):
            return sqrt(x ** 2 + y ** 2)
        
        def find_split(remain, mid):
            closer, far = [], []
            for i in remain:
                # equal must belong to closer, as we only do the update in closer.
                # e.g. [[0,1],[1,0]], 2. The mid: 0.5, 0.75, 0.875, 0.9375...
                if distance[i] <= mid: closer.append(i) 
                else: far.append(i)
            return (closer, far)
        
        remain = [i for i in range(len(points))]  # index
        distance = [dist(i, j) for i, j in points]
        left, right = 0, max(distance)
        closest = [] # res
        while k > 0:
            mid_dis = left + (right - left) / 2 # actual dist, not //2
            closer, far = find_split(remain, mid_dis)
            if len(closer) <= k:
                k -= len(closer)
                closest.extend(points[i] for i in closer)
                left = mid_dis
                remain = far
            else:
                right = mid_dis
                remain = closer
                
        return closest

“As the efficiency of this solution relies on averaging as close to a middle split of the points as possible on each iteration of the binary search, the use of Euclidean distances will be more efficient than the use of squared Euclidean distances. " The binary search solution works with squared distance as well. But on each iteration, the actual Euclidean distance will guarantee us to do the “Binary” search, which uses the exact middle distance to split the remaining points. Suppose the actual Euclidean distance is 1,2,3,4,5,6,7, we will use (1+7)/2 = 4 to split it evenly. But with squared distance, it’s 1,4,9,16.25,36,49, we will choose (1+49)/2 = 25 as the middle point to split it, it’s uneven, though we still can get a correct result.

  • code #quickSelection
class Solution:
    def kClosest(self, points: List[List[int]], k: int) -> List[List[int]]:
        def dis(x, y):
            return x ** 2 + y ** 2
        
        # provide a dis, pivot to split
        def partition(start, end):
            # split to points[:p], points[p:], first part < dis_split
            # dis_split = dis(*points[start]) can also use the first as split distance
            dis_split = dis(*points[start + (end - start) // 2]) # get a dist in between to split
            while start < end: # not <=, make sure swap is meaningful
                if dis(*points[start]) < dis_split:
                    start += 1
                else:
                    points[start], points[end] = points[end], points[start]
                    end -= 1
            # e.g. all nums <= split, start = end break the while, still needs to move start
            if dis(*points[start]) < dis_split: 
                start += 1
            return start # the start of second list
        
        start, end = 0, len(points) - 1
        p = len(points) # if equal, return directly
        while p != k:
            p = partition(start, end)
            if p < k:
                start = p
            else:
                end = p - 1
            
        return points[:k]  # k starts from 1