二分搜尋應用: 找第k小的配對距離_Find K-th Smallest Pair Dist_Leetcode #719

Akin Cakiner on Unsplash

題目敘述 719. Find K-th Smallest Pair Distance

給定一個輸入陣列nums和 參數k。

請找出第k小的pair distance是多少?

pair distance定義為 abs( nums[i] - nums[j]), i 不等於j

也就是任意兩陣列元素之間的差值絕對值。

測試範例

Example 1:

Input: nums = [1,3,1], k = 1
Output: 0
Explanation: Here are all the pairs:
(1,3) -> 2
(1,1) -> 0
(3,1) -> 2
Then the 1st smallest distance pair is (1,1), and its distance is 0.

第一小的pair distance 是 0​

Example 2:

Input: nums = [1,1,1], k = 2
Output: 0

第一小的pair distance 是 0
第二小的pair distance 是 0

Example 3:

Input: nums = [1,6,1], k = 3
Output: 5

第一小的pair distance 是 0
第二小的pair distance 是 5
第三小的pair distance 是 5

約束條件

Constraints:

• n == nums.length

n = 陣列長度。

• 2 <= n <= 10^4

陣列長度n介於2~一萬之間。

• 0 <= nums[i] <= 10^6

每個陣列元素都介於0~一百萬之間。

• 1 <= k <= n * (n - 1) / 2

參數k介於1 ~ C(n,2) 之間

觀察

除了暴力法計算全部pair dist再排序之外，其實還有更好的方法。

假如nums排序之後，那就變成從小到大依序放好的序列。

這時候可以用最小pair dist和最大pair dist作為上下區間，

往中間夾擠，看到哪一回合的時候，當下的pair dist剛好是第k小的。

演算法 排序前處理 + 二分搜尋法

承接觀察點的思考，先做前處理，將nums排序。

這樣就滿足binary search的前提，接著用最小pair dist和最大pair dist作為上下區間，

不斷砍半往中間夾擠pair dist，看到哪一回合的時候，當下的pair dist剛好是第k小的。

程式碼 排序前處理 + 二分搜尋法

class Solution(object):
    def countPairs(self, array, dist):
        count = 0

        # 累積看當下的 pair dist是第幾小的pair dist
        for i in range(len(array)):
            count += bisect.bisect_right(array, array[i] + dist, lo = i) - i - 1
        return count
        
    def smallestDistancePair(self, nums, k):
        """
        :type nums: List[int]
        :type k: int
        :rtype: int
        """
        nums.sort()
        low, high = min([abs(nums[i] - nums[i+1]) for i in range(len(nums) - 1)]), abs(nums[0] - nums[-1])

        while low < high:
            mid = (low + high)//2

            
            if self.countPairs(nums, mid) < k:
                # 假如不是第k小，而且還太小，則下一輪 pair dist猜大一點
                low = mid + 1
            elif self.countPairs(nums, mid) >= k:
                # 假如不是第k小，而且還太大，則下一輪 pair dist猜小一點
                # 假如mid是第k小，則下一輪 pair dist猜小一點，確認沒有更小的值也滿足
                high = mid
                
        return low

複雜度分析

時間複雜度: O( n log n)

前處理排序耗時O( n log n )

後續while loop + countPairs 總耗時也是O( n log n)

時間複雜度: O( n )

臨時變數用到長度為O(n)的list，所需空間為O(n)

關鍵知識點

遇到已排序的序列做搜尋，聯想到二分搜尋法。

如果輸入是亂序的，可以先做前處理排序，滿足二分搜尋法的前提。

強烈建議讀者一起復習相關的二分搜尋法框架和範例，鞏固知識點!

合縱連橫: 二分搜尋法框架_理解背後的本質

Reference

[1] 719. Find K-th Smallest Pair Distance