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]].
With Hoare partition scheme, the pivot and elements equal to the pivot can end up anywhere, and after a partition step p is not an index to the pivot, but just a separator, values to the left or at p are <= pivot, values to the right of p are >= pivot. With Hoare partition scheme quickselect requires recursing to the the base case of 1 element in order to find the kth element. If there are other elements equal to the kth element, they could end up on either side or both sides of the kth element.