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 on the plane. Find the closest points to the origin .

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


Solution :

Best Case QuickSelect

Approach 2: Divide and Conquer


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 and split the array into two buckets: one bucket of all the elements less than , and another bucket of all the elements greater than or equal to . This is known as "quickselecting by a pivot ".

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.


Let’s do the of partially sorting the subarray so that the smallest elements of this subarray occur in the first positions .

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

After, we have two buckets and , where are the original values when calling . Say the first bucket has items and the second bucket has items. If we were trying to partially sort say, items, then we only need to partially sort the first bucket: . Otherwise, if we were trying to partially sort say, items, then the first items are already partially sorted, and we only need to partially sort the next 7 items: .


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 .
  • Space Complexity: O(N).

Other approach:

Approach 1: Sort


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


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 .
  • Space Complexity: O(N).

Currently preparing for interviews

Get the Medium app

A button that says 'Download on the App Store', and if clicked it will lead you to the iOS App store
A button that says 'Get it on, Google Play', and if clicked it will lead you to the Google Play store