QuickSelect/QuickSort 如何从 Lomuto 分区方案转换为 Hoare 分区方案?

How to translate from Lomuto Partitioning scheme to Hoare's Partition scheme in QuickSelect/QuickSort?

我正在处理问题https://leetcode.com/problems/k-closest-points-to-origin/,问题陈述转载于此:

给定一个点数组,其中 points[i] = [xi, yi] 表示 X-Y 平面上的一个点和一个整数 k,return k 最近(欧氏距离)指向原点 (0, 0)k 个最近的点可以按任意顺序 return 编辑。

我正在使用 QuickSelect algorithm 来达到这个目的。这是我当前使用 Lomuto 分区方案(将最右边的元素作为基准)的有效但缓慢的代码。

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)

这是我用 Hoare 替换 Lomuto 的尝试。

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)       

但是,这个替换似乎出了问题。以下测试用例因我的 Hoare 尝试而失败:

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

我的输出是 [[-36,-57],[17,7],[-69,-8],[53,20],[-55,-58]],而预期输出是 [[17,7],[-2,-42],[53,20],[-36,-57],[-69,-8]]

使用 Hoare 分区方案,主元和等于主元的元素可以在任何地方结束,并且在分区步骤之后 p 不是主元的索引,而只是一个分隔符,值在左边或者在 p 处 <= pivot,p 右侧的值 >= pivot。使用 Hoare 分区方案,quickselect 需要递归到 1 个元素的基本情况才能找到第 k 个元素。如果还有其他元素等于第k个元素,它们可能在第k个元素的任一侧或两侧结束。