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

2563. 统计公平数对的数目

题目要统计和落在 [lower, upper] 内的数对。统计的是不同位置之间的配对数量，重新排列元素不会改变这个数量，所以可以先排序，再利用二分查找确定范围。

由 $lower\le x+y\le upper$ 可以得到 $lower-x\le y \le upper-x$ 。对于位置 i 上的 x，只在 i+1 到 n-1 的后缀中查找，避免同一对被重复统计。Python 中的 bisect 返回右侧插入位置，因此 L=bisect(a,lower-x-1)，R=bisect(a,upper-x)-1，两次查找都要传入起点 i+1。闭区间 [L, R] 内有 R-L+1 个符合条件的元素。

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 返回迭代器。要计算数组下标，可以将它减去 nums.begin()，再做下标的加减；如果只需要统计区间长度，也可以直接计算两个边界迭代器之差。