LeetCode719. 找出第 K 小的数对距离

题目描述

本题目来自LeetCode上的『719. 找出第 K 小的数对距离』

数对 (a,b) 由整数 a 和 b 组成，其数对距离定义为 a 和 b 的绝对差值。

给你一个整数数组 nums 和一个整数 k ，数对由 nums[i] 和 nums[j] 组成且满足 0 <= i < j < nums.length 。返回 所有数对距离中 第 k 小的数对距离。

示例1：

输入：nums = [1,3,1], k = 1
输出：0
解释：数对和对应的距离如下：
(1,3) -> 2
(1,1) -> 0
(3,1) -> 2
距离第 1 小的数对是 (1,1) ，距离为 0 。

提示

n == nums.length
2 <= n <= 10^4
0 <= nums[i] <= 10^6
1 <= k <= n * (n - 1) / 2

题解

设 $m=max(nums)-min(nums)$，所有 数对距离 按升序排序为 $a_1,a_2,…,a_k,…,a_m$，其中 $a_k$ 为第 $k$ 小的数。通过枚举 $a_i$ 的值，计算 $a_i$ 是第几小的数，即可得到答案。

枚举可以用二分查找替代，查找范围为 $[min(nums),max(nums)]$，若 $mid$ 在 $a_k$ 的左边，则下次查找的范围为 $[mid + 1, right]$，若 $mid$ 在 $a_k$ 的右边（包含 $a_k$ 的位置），则下次查找的范围为 $[left, mid]$。

计算 $a_i$ 是第几小的数可以使用双指针，左端点设为 $i$，枚举右端点 $j$ 直至不符合条件，则 $j$ 是第一个不符合条件的位置，因此右端点的位置范围为 $[i+1,j-1]$，共 $j-i-1$ 个。

代码

class Solution {
public:
    int smallestDistancePair(vector<int>& nums, int k) {
        int n = nums.size();
        sort(nums.begin(), nums.end());
        int left = 0, right = nums.back() - nums.front();

        auto check = [&](int x) -> int {
            int res = 0;
            for (int i = 0, j = 1; i < n; ++i) {
                while (j < n && nums[j] - nums[i] <= x) {
                    ++j;
                }
                res += j - i - 1;
            }
            return res;
        };

        while (left < right) {
            int mid = left + ((right - left) >> 1);
            if (check(mid) >= k) {
                right = mid;
            } else {
                left = mid + 1;
            }
        }
        return left;
    }
};