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

題目連結：719. Find K-th Smallest Pair Distance

題目敘述

#

The distance of a pair of integers a and b is defined as the absolute difference between a and b.

Given an integer array nums and an integer k, return the k th smallest distance among all the pairs nums[i] and nums[j] where 0 <= i < j < nums.length.

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.

Example 2:

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

Example 3:

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

解題思路

#

由於這題 nums 的大小 n 達到 10^4，因此如果我們直接窮舉所有的 pair 肯定會 TLE，我們需要使用更有效率的方法。

這題 nums[i] 最大值只有 10^6，我們可以使用 Binary Search，來猜測一個可能的距離 x，並計算距離小於等於 x 的 pair 數有多少：

• 如果 pair 數太少，我們就往更大的 x 去猜。
• 如果 pair 數太多，我們就往更小的 x 去猜。
• 如果 pair 數等於 k，這時答案還不一定會是 x，我們有可能還是猜得太大了。可以先把這個 x 存起來，再往更小的 x 去猜。

這時問題變成：我們該如何知道距離小於等於 x 的 pair 數有多少呢？

我們可以先將 nums 從小到大做排序。

這時可以發現，假設給定兩個 index：l 與 r，在 l 到 r 之間的所有元素的距離，絕對不會超過 nums[r] - nums[l]。 所以我們可以使用 Sliding Window 的想法，維護兩個 pointer，來計算有多少個 pair 距離小於等於 x：

• 一開始先初始化 l = 0，r = 1。
• 如果 nums[r] - nums[l] <= x，則 l 到 r 之間的所有元素的距離都不會超過 x，且相較於 l 到 r-1，新增的 pair 數量為 r - l。將 r - l 累加到 cnt 上，並嘗試將 r 向右移動。
• 如果 nums[r] - nums[l] > x，則嘗試將 l 往右移動，直到 nums[r] - nums[l] <= x。
• 最後 cnt 就會是距離小於等於 x 的 pair 數。

程式碼

#

C++

#

class Solution {
public:
    int dis(int a, int b) {
        return abs(a - b);
    }

    int check(int x, vector<int>& nums, int k) {
        // sliding windows
        int cnt = 0;
        int l = 0, r = 1;
        while (r < nums.size()) {
            if (dis(nums[l], nums[r]) <= x) {
                cnt += (r - l);
                r++;
            } else {
                l++;
            }
        }
        return cnt;
    }

    int binarySearch(vector<int>& nums, int k) {
        int l = 0;
        int r = *(nums.end() - 1); // Last element
        
        while (l <= r) {
            int m = (l + r) >> 1;

            //  Check how many pairs have distance <= m
            if (check(m, nums, k) < k) {
                l = m + 1; // 嘗試往更大的 m 搜尋
            } else {
                r  = m - 1; // 嘗試往更小的 m 搜尋
            }
        }
        return l;
    }

    int smallestDistancePair(vector<int>& nums, int k) {
        sort(nums.begin(), nums.end());
        return binarySearch(nums, k);
    }
};