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

Difficulty
Hard
Tags
二分查找
URL
https://leetcode.cn/problems/find-k-th-smallest-pair-distance/
Star
数对 (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 。
示例 2：
输入：nums = [1,1,1], k = 2
输出：0
示例 3：
输入：nums = [1,6,1], k = 3
输出：5
提示：
  • n == nums.length
  • 2 <= n <= 104
  • 0 <= nums[i] <= 106
  • 1 <= k <= n * (n - 1) / 2
通过次数22,648提交次数51,479

法 1 二分法

思路
在看到k值问题时我们比较容易想到堆，再一个便是二分。我们先来考虑一下堆，堆在本题应该只有一种使用方式，便是用两个for loop循环将所有可能的数对距离加入到堆中，在将堆pop k 次。以上算法时间复杂度明显大于O(n^2)。毫无技巧可言，属于暴力解法，不出意外会超时。因此考虑使用二分。
二分思路：首先我们不考虑细节问题，只考虑一下二分简单的逻辑应该怎么找到这个k-th距离。我们会想到的是有一个左边界与一个右边界，每次都不断地缩小这两个边界。如何缩小呢？我们找到两个边界的中间值mid = (left + right)/2，我们来看看当前的以此mid为距离时，这个距离在数组中有多少对数对比它小，如果有x个数对比此中间值小（小于等于)，那么是否说明此中间值mid就是第x+1大的数对距离。
因此我们便尝试去统计每一次该中间值mid是第多少对数对距离记作cnt，并且与k比较:
  • 如果与cnt < k，则说明mid此时的距离不够，需要增大距离靠近k。
  • 反正cnt >= k，则说明mid此时的距离有可能刚刚好是在第k个，或者太大了，则需要缩小距离靠近k。
如何去计算该数组中以此mid为距离时，这个距离在数组中有多少对数对比它小？朴素的做法就是每一次都进行O(n^2)的遍历然后加一统计，并且每一次二分时都需要这样。
更为快速则做法则是(排序+滑动窗口)。
题解
class Solution:
    def smallestDistancePair(self, nums: List[int], k: int) -> int:
        nums.sort()
        size = len(nums)

        l, r = 0, nums[-1] - nums[0]
        while l <= r:
            mid = l + (r - l) // 2
            count = self.getCount(mid, nums, size)
            if count < k:
                l = mid + 1
            elif count >= k:
                r = mid - 1

        return l

    # 统计一下距离 小于等于 d 的个数, 有序数组,直接上滑动窗口
    def getCount(self, d, nums, size):
        count = 0
        l, r = 0, 1
        # 以 nums[r] 为右端点，不断的收缩左端点
        while r < len(nums):
            cur_val = nums[r]
            while l < r and nums[r] - nums[l] > d:
                l += 1
            count += r - l 
            r += 1
            
        return count