0719. 找出第 k 小的距离对 #

标签：数组、双指针、二分查找、排序
难度：困难

题目大意 #

给定一个整数数组 nums，对于数组中不同的数 nums[i]、nums[j] 之间的距离定义为 nums[i] 和 nums[j] 的绝对差值，即 dist(nums[i], nums[j]) = abs(nums[i] - nums[j])。求所有数对之间第 k 个最小距离。

解题思路 #

一般来说 topK 问题都可以用堆排序来解决。但是这道题使用堆排序超时了。所以需要换其他方法。

先来考虑第 k 个最小距离的范围。这个范围一定在 [0, max(nums) - min(nums)] 之间。

我们可以对 nums 先进行排序，然后得到最小距离为 0，最大距离为 nums[-1] - nums[0]。我们可以在这个区间上进行二分，对于二分的位置 mid，统计距离小于等于 mid 的距离对数，并根据它和 k 的关系调整区间上下界。

统计对数可以使用双指针来计算出所有小于等于 mid 的距离对数目。维护两个指针 left、right。left、right 都指向数组开头位置。然后不断移动 right，计算 nums[right] 和 nums[left] 之间的距离，如果大于 mid，则 left 向右移动，直到距离小于等于 mid 时，统计当前距离对数为 right - left。最终将这些符合要求的距离对数累加，就得到了所有小于等于 mid 的距离对数目。

代码 #

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19


class Solution:
    def smallestDistancePair(self, nums: List[int], k: int) -> int:
        def get_count(dist):
            left, count = 0, 0
            for right in range(1, len(nums)):
                while nums[right] - nums[left] > dist:
                    left += 1
                count += (right - left)
            return count

        nums.sort()
        left, right = 0, nums[-1] - nums[0]
        while left < right:
            mid = left + (right - left) // 2
            if get_count(mid) >= k:
                right = mid
            else:
                left = mid + 1
        return left