K Closest Points to Origin — Leetcode Medium

973. K Closest Points to Origin

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

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

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]]
(The answer [[-2,4],[3,3]] would also be accepted.)

Note:

1. 1 <= K <= points.length <= 10000
2. -10000 < points[i][0] < 10000
3. -10000 < points[i][1] < 10000

Solution :

Best Case QuickSelect

Approach 2: Divide and Conquer

Intuition

We want an algorithm faster than Nlog⁡NNlogN. Clearly, the only way to do this is to use the fact that the K elements returned can be in any order — otherwise we would be sorting which is at least Nlog⁡NNlogN.

Say we choose some random element x = A[i] and split the array into two buckets: one bucket of all the elements less than x, and another bucket of all the elements greater than or equal to x. This is known as "quickselecting by a pivot x".

The idea is that if we quickselect by some pivot, on average in linear time we’ll reduce the problem to a problem of half the size.

Algorithm

Let’s do the work(i, j, K) of partially sorting the subarray (points[i], points[i+1], ..., points[j]) so that the smallest K elements of this subarray occur in the first K positions (i, i+1, ..., i+K-1).

First, we quickselect by a random pivot element from the subarray. To do this in place, we have two pointers i and j, and move these pointers to the elements that are in the wrong bucket -- then, we swap these elements.

After, we have two buckets [oi, i] and [i+1, oj], where (oi, oj) are the original (i, j) values when calling work(i, j, K). Say the first bucket has 10items and the second bucket has 15 items. If we were trying to partially sort say, K = 5 items, then we only need to partially sort the first bucket: work(oi, i, 5). Otherwise, if we were trying to partially sort say, K = 17 items, then the first 10items are already partially sorted, and we only need to partially sort the next 7 items: work(i+1, oj, 7).

Solution:

We implement partition and then we also implement the Sort()

def kClosest(self, points, K):
    dist = lambda i: points[i][0]**2 + points[i][1]**2
    # Quick sort partitioning using Horre's partitioning
    def partition(i, j):
        # Partition by pivot A[i], returning an index mid
        # such that A[i] <= A[mid] <= A[j] for i < mid < j.
        pi = i
        pivot = dist(i)
        i += 1
        while True:
            while i < j and dist(i) < pivot:
                i += 1
            while i <= j and dist(j) >= pivot:
                j -= 1
            if i >= j: break
            # swap where the dist(i) > pivot and dist(j) < pivot
            points[i], points[j] = points[j], points[i]
# Myself: Now all points upto the jth are smaller than pivot
        # move the pivot to j and return pivot's index which is j.
        points[pi], points[j] = points[j], points[pi]
        
        return j
    
    
    def sort(i, j, K):
        # Partially sorts A[i:j+1] so the first K elements are
        # the smallest K elements.
        if i >= j: return
        # Put random element as A[i] - this is the pivot
        k = random.randint(i, j)
        points[i], points[k] = points[k], points[i]
        pi = partition(i, j)
        if K < pi - i + 1:
             sort(i, pi - 1, K)
        elif K > pi - i + 1:
             # Notice how we reduced K to K - (pi - i +1)
             sort(pi + 1, j, K - (pi - i + 1))
    sort(0, len(points) - 1, K)
    # our purpose is the ensure all k closest elements are
    # included in the first K positions. So We do QuickSelect till out pivot is at the Kth position. Then we quit.
    return points[:K]

Complexity Analysis

• Time Complexity: O(N) in average case and O(N2) in the worst case, where N is the length of points.
• Space Complexity: O(N).

Other approach:

Approach 1: Sort

Intuition

Sort the points by distance, then take the closest K points.

Algorithm:

In Python, we sort by a custom key function — namely, the distance to the origin. Afterwards, we return the first K elements of the list.

class Solution(object):
    def kClosest(self, points, K):
        points.sort(key = lambda P: P[0]**2 + P[1]**2)
        return points[:K]

Complexity Analysis

• Time Complexity: O(Nlog⁡N), where N is the length of points.
• Space Complexity: O(N).