如何在 QuickSelect/QuickSort 中從 Lomuto 分區方案轉換為 Hoare 的分區方案？

Question

我正在解決問題https://leetcode.com/problems/k-closest-points-to-origin/此處轉載了問題聲明：

給定一個點數組，其中points[i] = [xi, yi]表示 XY 平面上的一個點和一個整數k ，返回離原點(0, 0)最接近（歐幾里得距離）的k個點。 可以按任何順序返回k個最近點。

為此，我正在使用QuickSelect 算法。 這是我當前使用 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]] 。

Answer 1

使用 Hoare 分區方案，樞軸和等於樞軸的元素可以在任何地方結束，並且在分區步驟之后p不是樞軸的索引，而只是一個分隔符，左側或p處的值是 <= 樞軸，值p的右側是 >= 樞軸。 使用 Hoare 分區方案，快速選擇需要遞歸到 1 個元素的基本情況才能找到第 k 個元素。 如果有其他元素等於第 k 個元素，它們可能會出現在第 k 個元素的任一側或兩側。