🔗 🔴 719. Find K-th Smallest Pair Distance

tags: Weekly Contest 56 陣列(Array) 排序(Sorting) 二分搜尋(Binary Search) 雙指針(Two Pointers) 滑動窗口(Sliding Window)

題意

一對整數 (a,b)(a, b) 由整數 aa 和 bb 組成，其距離定義為 aa 和 bb 之間的絕對差 ∣a−b∣|a - b|。

給定一個整數陣列 numsnums 和一個整數 kk，返回所有 nums[i]nums[i] 和 nums[j]nums[j] 配對中第 kk 小的距離，其中 0≤i<j<nums.length0 \leq i < j < nums.length。

約束條件

  • n=nums.lengthn = nums.length
  • 2≤n≤1042 \leq n \leq 10^4
  • 0≤nums[i]≤1060 \leq nums[i] \leq 10^6
  • 1≤k≤n∗(n−1)21 \leq k \leq \frac{n * (n - 1)}{2}

為方便敘述，以下標配對 (i,j)(i, j) 表示 nums[i]nums[i] 和 nums[j]nums[j] 的配對，即 (nums[i],nums[j])(nums[i], nums[j]) 。

雖然題目中的數對 (i,j)(i, j) 需要滿足 i<ji < j，但注意到距離的定義是 絕對值 ，因此從原陣列中找數對 (i,j)(i, j) ，等同於在 排序 陣列中找數對 (i,j)(i, j)。

最暴力的做法是根據所有可能的 (i,j)(i, j) 的配對，生成所有距離，然後排序後取出第 kk 小的距離即可。但所有可能的配對有 n∗(n−1)2\frac{n * (n - 1)}{2} 個，故暴力法的時間複雜度為 O(n2log⁡n2)O(n^2 \log n^2) 。而在約束條件中 nn 為 10410^4 ，顯然是會超時的。

如果我們將所有數對的距離進行 計數(Counting) ，那我們可以得到一個距離和數量的陣列 cntcnt，其中 cnt[i]cnt[i] 表示所有 (i,j)(i, j) 數對中，距離為 ii 的數對數量。接著對 cntcnt 做 前綴和(Prefix Sum) 得到 ss，則 s[k]s[k] 表示 距離 ≤k\leq k 的數對數量。由於 cnt[i]cnt[i] 皆為非負整數，因此 ss 具有 遞增 性質。

而若 s[k]=y,s[k−1]=xs[k] = y, s[k-1] = x，則代表第 x+1,x+2,...,yx+1, x+2, ..., y 小的距離皆為 kk ，這裡可以代範例下去理解。

因此若可以在不生成所有數對的情況下，就可以確定 s[k]s[k] 的值，就可以對所有可能的距離做 二分搜尋(Binary Search) 。又數對的距離最小為 00 ，最大為 U−LU - L，其中 UU 和 LL 分別為整數陣列的最大值和最小值，因此二分的範圍為 [0,U−L][0, U-L] 。

令 check(k) 表示 s[k]s[k] ，即距離 ≤k\leq k 的數對數量，則有兩種方式可以直接確定 s[k]s[k] 的值：

  1. 二分搜尋(Binary Search)
    • 枚舉所有可能的右端點 jj，使用二分搜尋找出左端點 ii，使得數對 (i,j),(i+1,j),...,(j−1,j)(i, j), (i+1, j), ..., (j-1, j) 的距離都小於等於 kk 。
  2. 滑動窗口(Sliding Window)
    • 同樣枚舉所有可能的右端點 jj，但不用重新計算，而是逐漸縮小左端點 ii，使得數對 (i,j),(i+1,j),...,(j−1,j)(i, j), (i+1, j), ..., (j-1, j) 的距離都小於等於 kk 。

在確定右端點 jj 之後，若要使得數對 (i,j),(i+1,j),...,(j−1,j)(i, j), (i+1, j), ..., (j-1, j) 這 j−ij - i 個數對的距離都小於等於 kk ，則 nums[i]nums[i] 的值需要至少為 k−nums[j]k - nums[j] ，因此可以使用 bisect_left 或 lower_bound 來找出左端點 ii ，即第一個滿足 nums[i]≥k−nums[j]nums[i] \geq k - nums[j] 的 ii 。

枚舉所有可能的 jj，累加所有滿足條件的數對數量即可。

複雜度分析

  • 時間複雜度：O(nlog⁡n×log⁡(U−L))\mathcal{O}(n \log n \times \log (U - L))，其中 nn 為陣列的長度，UU 和 LL 分別為整數陣列的最大值和最小值。
    • 排序陣列需要 O(nlog⁡n)O(n \log n) 的時間。
    • 對答案做二分搜尋需要 O(log⁡(U−L))O(\log (U - L)) 的時間。
    • 計算 check(k) 需要 O(nlog⁡n)O(n \log n) 的時間。
  • 空間複雜度：O(log⁡n)\mathcal{O}(\log n)，即排序陣列的空間複雜度。
1
2
3
4
5
6
7
8
9
10
11
12
13
14
class Solution:
    def smallestDistancePair(self, nums: List[int], k: int) -> int:
        n = len(nums)
        nums.sort() # 由小到大排序
        
        def check(mid: int) -> int: # 計算距離 <= mid 的 pair 數量
            res = 0
            for j, x in enumerate(nums): # 枚舉右端點 j
                i = bisect_left(nums, x - mid, 0, j)
                res += j - i # (i, j), (i+1, j), ..., (j-1, j) 距離都小於等於 mid
            return res
        
        # 對答案做二分搜尋
        return bisect_left(range(nums[n-1] - nums[0]), k, key=check)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
class Solution {
public:
    int smallestDistancePair(vector<int>& nums, int k) {
        int n = nums.size();
        sort(nums.begin(), nums.end()); // 由小到大排序

        function<int(int)> check = [&](int mid) { // 計算距離 <= mid 的 pair 數量
            int res = 0;
            for (int j = 0; j < n; j++) { // 枚舉右端點 j
                int i = lower_bound(nums.begin(), nums.begin() + j, nums[j] - mid) - nums.begin();
                res += j - i; // (i, j), (i+1, j), ..., (j-1, j) 距離都小於等於 mid
            }
            return res;
        };

        // 對答案做二分搜尋
        int left = 0, right = nums[n-1] - nums[0]; 
        while (left <= right) {
            int mid = left + (right - left) / 2;
            if (check(mid) >= k) right = mid - 1;
            else left = mid + 1;
        }
        return left;
    }
};

方法二：滑動窗口(Sliding Window)

由於陣列已經排序過，因此 (i,j)(i, j) 的距離會比 (i,j−1)(i, j - 1) 大，故若 (i,j)(i, j) 的不滿足條件(即距離大於 kk)，則(i - 1, j) 的距離只會更大，顯然也不滿足條件。因此我們實際上不需要重新計算左端點 ii，只需要檢查先前的左端點 ii 和當前的右端點 jj 是否符合條件即可。若不滿足條件，再移動左端點 ii 直到滿足條件即可。

雖然看似是 O(n2)O(n^2) 的時間複雜度，但實際上只需要 O(n)O(n) 的時間複雜度，因為左端點 ii 最多只會從 00 走到 n−1n-1 。

複雜度分析

  • 時間複雜度：O(nlog⁡n+nlog⁡(U−L))\mathcal{O}(n \log n + n \log (U - L))，其中 nn 為陣列的長度，UU 和 LL 分別為整數陣列的最大值和最小值。
    • 排序陣列需要 O(nlog⁡n)O(n \log n) 的時間。
    • 對答案做二分搜尋需要 O(log⁡(U−L))O(\log (U - L)) 的時間。
    • 計算 check(k) 需要 O(n)O(n) 的時間。
  • 空間複雜度：O(log⁡n)\mathcal{O}(\log n)，即排序陣列的空間複雜度。
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
class Solution:
    def smallestDistancePair(self, nums: List[int], k: int) -> int:
        n = len(nums)
        nums.sort() # 由小到大排序

        def check(mid: int) -> int: # 計算距離 <= mid 的 pair 數量
            res = left = 0
            for right, x in enumerate(nums): # 枚舉右端點 right
                while x - nums[left] > mid: # 移動左端點直到滿足條件為止
                    left += 1
                res += right - left # (left, right), (left+1, right), ..., (right-1, right) 距離都小於等於 mid
            return res
        
        # 對答案做二分搜尋
        return bisect_left(range(nums[n-1] - nums[0]), k, key=check)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
class Solution {
public:
    int smallestDistancePair(vector<int>& nums, int k) {
        int n = nums.size();
        sort(nums.begin(), nums.end()); // 由小到大排序

        function<int(int)> check = [&](int mid) { // 計算距離 <= mid 的 pair 數量
            int res = 0, l = 0;
            for (int r = 0; r < n; r++) { // 枚舉右端點 r
                while (nums[r] - nums[l] > mid) { // 移動左端點直到滿足條件為止
                    l++;
                }
                res += r - l; // (l, r), (l+1, r), ..., (r-1, r) 距離都小於等於 mid
            }
            return res;
        };

        // 對答案做二分搜尋
        int left = 0, right = nums[n-1] - nums[0]; 
        while (left <= right) {
            int mid = left + (right - left) / 2;
            if (check(mid) >= k) right = mid - 1;
            else left = mid + 1;
        }
        return left;
    }
};

寫在最後

masterpiece, best quality, high quality,extremely detailed CG unity 8k wallpaper, extremely detailed, High Detail, colors,
(1girl, solo), (idol, idol costume), long hair, black hair, dress, bow, standing, detached sleeves, white dress, hand on hip, curtains, pointing, pointing at self, stage, on stage,
A young girl wearing a lavish purple dress with puffy sleeves and a layered skirt, Her hair is styled in twin tails with purple bows, The background is a dark blue curtain, She is smiling and posing cutely,

雖然程式碼可能就短短幾行，但寫 Hard 的解題紀錄除了相比於 Easy 和 Medium 需要花更多的時間外，也需要注意到更多東西。但在完整闡述思考過程後，還是有助於自己釐清思路。

因此雖然很忙，還是要寫一些 Hard 的解題紀錄，不然感覺腦袋都變遲鈍了。