I am working on the problem https://leetcode.com/problems/k-closest-points-to-origin/ with the question statement reproduced here:
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 (Euclidean distance) points to the origin (0, 0)
. The k
closest points may be returned in any order.
I am using the QuickSelect algorithm for this purpose. Here is my current, working but slow code using the Lomuto partitioning scheme (taking the rightmost element to be the pivot).
class Solution:
def kClosest(self, points: List[List[int]], k: int) -> List[List[int]]:
# attempt at using quickselect
def dist(point):
a, b = point
return a ** 2 + b ** 2
def quickSelect(l, r):
# Using the (slower) Lomuto Partioning Scheme
pivot, p = points[r], l
for i in range(l, r):
if dist(points[i]) <= dist(pivot):
points[p], points[i] = points[i], points[p] # swap them
p += 1
points[p], points[r] = points[r], points[p]
# if the pointer's index is greater than the desired index k,
# then we need to narrow the range
if p == k - 1: return points[:k]
elif p < k - 1: return quickSelect(p + 1, r)
else: return quickSelect(l, p - 1)
return quickSelect(0, len(points) - 1)
Here is my attempt at replacing Lomuto with Hoare.
class Solution:
def kClosest(self, points: List[List[int]], k: int) -> List[List[int]]:
# attempt at using quickselect
def dist(point):
a, b = point
return a ** 2 + b ** 2
def quickSelect(l, r):
# Using the (faster) Hoare scheme
pivot_index = ((r + l) // 2)
pivot = points[pivot_index]
i, j = l - 1, r + 1
while True:
while True:
i += 1
if dist(points[i]) >= dist(pivot):
break
while True:
j -= 1
if dist(points[j]) <= dist(pivot):
break
if i >= j:
p = j
break
points[i], points[j] = points[j], points[i]
# if the pointer's index is greater than the desired index k,
# then we need to narrow the range
if p == k - 1: return points[:k]
elif p < k - 1: return quickSelect(p + 1, r)
else: return quickSelect(l, p - 1)
return quickSelect(0, len(points) - 1)
However, it seems as though this replacement has gone awry. The following test case is failing with my Hoare attempt:
points = [[-95,76],[17,7],[-55,-58],[53,20],[-69,-8],[-57,87],[-2,-42],[-10,-87],[-36,-57],[97,-39],[97,49]]
5
k = 5
My output is [[-36,-57],[17,7],[-69,-8],[53,20],[-55,-58]]
while the expected output is [[17,7],[-2,-42],[53,20],[-36,-57],[-69,-8]]
.