如何在QuickSelect/QuickSort中将Lomuto分区方案转换为Hoare的分区方案？

Question

我正在研究问题https://leetcode.com/problems/k-closest-points-to-origin/，并在此转载问题陈述：

给定一个点数组，其中points[i] = [xi, yi]表示X平面上的一个点和一个整数k，则将k最近的(欧氏距离)点返回到原点(0, 0)。k最近点可以按任何顺序返回。

我正为此目的使用QuickSelect算法。下面是我目前正在工作但速度较慢的代码，它使用的是洛穆托分区方案(将最右边的元素作为支点)。

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]]。

Answer 1

使用Hoare分区方案，枢轴和与支点相等的元素可以在任何地方结束，在分区步骤p不是枢轴的索引，而只是分隔符，左边或p的值是<=支点，p右边的值是>=支点。使用Hoare分区方案，quickselect需要递归到1元素的基本情况，以便找到kth元素。如果有与kth元素相等的其他元素，则它们可能在kth元素的任何一侧或两侧结束。