leetcode-2563. 统计公平数对的数目

LeetCode 2563. 统计公平数对的数目

1. 解题思路

题目要求统计数组中满足以下条件的数对数量：

也就是说，我们需要找出两个不同位置的元素，使它们的和落在 [lower, upper] 区间内。

最直接的方法是使用双重循环，枚举所有可能的数对，逐个判断是否满足条件。不过，这样的时间复杂度为 $O(n^2)$，当数组较长时，效率就不太理想了。

仔细想想，题目只要求统计符合条件的数对数量，并不关心它们原本在数组中的位置。因此，我们可以先对数组排序，再利用二分查找快速确定符合条件的元素范围。

2. 排序与二分查找

假设当前遍历到的元素为 x，另一个元素为 y。

根据题目的要求，有：

将 x 移到不等式右边：

这样，问题就变成了：对于当前元素 x，在数组中找到所有处于 [lower - x, upper - x] 区间内的元素，并统计它们的数量。

由于数组已经排好序，我们可以通过两次二分查找，分别确定这个区间的左右边界。

不过，还有一个细节需要注意。

题目要求 i < j，因此对于位置 i 上的元素 x，我们只需要在它后面的区间 [i + 1, n - 1] 中查找。

这样既能保证不会将同一个位置的元素使用两次，也能避免同一对元素被重复统计。

3. 使用 bisect 确定边界

Python 的 bisect 模块提供了二分查找功能。其中，bisect() 默认等价于 bisect_right()，返回目标值在有序数组中的右侧插入位置。

这里需要找到两个边界：

• L：第一个大于等于 lower - x 的位置。
• R：最后一个小于等于 upper - x 的位置。

由于 bisect() 返回的是右侧插入位置，我们可以稍微调整查找的目标值。

左边界：

L = bisect(a, lower - x - 1, i + 1)

因为题目中的元素都是整数，所以第一个大于 lower - x - 1 的元素，也就是第一个大于等于 lower - x 的元素。

右边界：

R = bisect(a, upper - x, i + 1) - 1

bisect() 会返回最后一个小于等于 upper - x 的元素之后的位置，因此再减去 1，就得到了我们需要的右边界。

注意，两次查找都传入了 i + 1，将搜索范围限制在当前元素后面的部分。

找到边界后，闭区间 [L, R] 中符合条件的元素数量就是：

如果不存在符合条件的元素，计算结果也会是 0，因此不需要额外判断。

将每次查询得到的数量累加起来，就得到了最终答案。

4. 代码实现

Python 代码：

from bisect import *

class Solution:
    def countFairPairs(self, a: List[int], lower: int, upper: int) -> int:
        a.sort()
        res = 0
        for i, x in enumerate(a):
            L = bisect(a, lower - x - 1, i + 1)
            R = bisect(a, upper - x, i + 1) - 1
            res += R - L + 1
        return res

C++ 代码：

class Solution {
public:
    long long countFairPairs(vector<int>& nums, int lower, int upper) {
        sort(nums.begin(), nums.end());
        long long res = 0;
        for (int i = 0; i < nums.size(); i++) {
            int L = upper_bound(nums.begin() + i + 1, nums.end(),
                                lower - nums[i] - 1) -
                    nums.begin();
            int R =
                upper_bound(nums.begin() + i + 1, nums.end(), upper - nums[i]) -
                nums.begin() - 1;
            res += R - L + 1;
        }
        return res;
    }
};

C++ 中的 upper_bound() 与 Python 的 bisect() 类似，都会查找目标值的右侧插入位置。

不过，两者的返回值有所不同：

• Python 的 bisect() 直接返回数组下标。
• C++ 的 upper_bound() 返回迭代器，需要减去 nums.begin() 才能得到对应的数组下标。

因此，C++ 代码中的：

upper_bound(nums.begin() + i + 1, nums.end(), upper - nums[i])
    - nums.begin() - 1

就是先取得右侧插入位置，再转换为数组下标，最后减去 1，得到符合条件的右边界。

如果只是为了统计元素数量，也可以直接计算两个边界迭代器之间的距离，不必先转换成下标。

5. 复杂度分析

排序需要 $O(n\log n)$ 的时间。

之后遍历数组中的每一个元素，并进行两次二分查找。每次二分查找的时间复杂度为 $O(\log n)$，因此这一部分的总时间复杂度也是 $O(n\log n)$。

所以，整个算法的时间复杂度为：

相比直接枚举所有数对的 $O(n^2)$，效率有了明显提升。

这道题的关键是将数对的和满足某个区间条件，转化成固定一个元素后，另一个元素需要落在哪个区间内。

完成这个转换后，就可以利用有序数组和二分查找快速统计答案，而不必枚举每一组数对。