满足不等式的数对数目

难度: Hard

给你两个下标从 0 开始的整数数组 nums1 和 nums2 ，两个数组的大小都为 n ，同时给你一个整数 diff ，统计满足以下条件的数对 (i, j) ：

0 <= i < j <= n - 1 且
nums1[i] - nums1[j] <= nums2[i] - nums2[j] + diff.

请你返回满足条件的 数对数目 。

示例 1：

输入：nums1 = [3,2,5], nums2 = [2,2,1], diff = 1
输出：3
解释：
总共有 3 个满足条件的数对：
1. i = 0, j = 1：3 - 2 <= 2 - 2 + 1 。因为 i < j 且 1 <= 1 ，这个数对满足条件。
2. i = 0, j = 2：3 - 5 <= 2 - 1 + 1 。因为 i < j 且 -2 <= 2 ，这个数对满足条件。
3. i = 1, j = 2：2 - 5 <= 2 - 1 + 1 。因为 i < j 且 -3 <= 2 ，这个数对满足条件。
所以，我们返回 3 。

示例 2：

输入：nums1 = [3,-1], nums2 = [-2,2], diff = -1
输出：0
解释：
没有满足条件的任何数对，所以我们返回 0 。

提示：

n == nums1.length == nums2.length
2 <= n <= 10⁵
-10⁴ <= nums1[i], nums2[i] <= 10⁴
-10⁴ <= diff <= 10⁴

Submission

运行时间: 306 ms

内存: 33.9 MB

class Solution:
    def numberOfPairs(self, nums1: List[int], nums2: List[int], diff: int) -> int:
        nums=[i-j for i,j in zip(nums1,nums2)]
        l=list(set(nums))
        l.sort()
        cnt=Counter()
        for i in range(len(l)):
            cnt[l[i]]=i
        class Tree:
            def __init__(self,num):
                n=len(num)
                self.tree=[0]*(n+1)
            def update(self,i):
                i+=1
                while i<len(self.tree):
                    self.tree[i]+=1
                    i+=i&-i
            def presum(self,i):
                i+=1
                s=0
                while i:
                    s+=self.tree[i]
                    i-=i&-i
                return s
        tr=Tree(l)
        ans=0
        for k in nums:
            t=k+diff
            ind=bisect_right(l,t)
            ans+=tr.presum(ind-1)
            tr.update(cnt[k])
        return ans

Explain

这道题的题解采用了树状数组 (Binary Indexed Tree, BIT) 的数据结构来优化时间复杂度。首先，计算 nums1 和 nums2 对应元素的差值，存储在数组 nums 中。然后，对 nums 中的元素去重并排序，得到数组 l。接着，使用一个哈希表 cnt 来记录每个元素在数组 l 中的索引。通过树状数组 tr 来维护前缀和，遍历 nums 数组，对于每个元素 k，找到 k+diff 在数组 l 中的右边界索引 ind（即第一个大于 k+diff 的元素的索引），通过树状数组查询 [0, ind-1] 范围内的元素个数，并累加到答案 ans 中。最后，更新树状数组，将当前元素 k 对应的索引在树状数组中的值加 1。

时间复杂度: O(nlogn)

空间复杂度: O(n)

class Solution:
    def numberOfPairs(self, nums1: List[int], nums2: List[int], diff: int) -> int:
        nums = [i - j for i, j in zip(nums1, nums2)]  # 计算差值
        l = list(set(nums))  # 去重
        l.sort()  # 排序
        cnt = Counter()  # 哈希表记录每个元素的索引
        for i in range(len(l)):
            cnt[l[i]] = i
        class Tree:
            def __init__(self, num):
                n = len(num)
                self.tree = [0] * (n + 1)
            def update(self, i):
                i += 1
                while i < len(self.tree):
                    self.tree[i] += 1
                    i += i & -i
            def presum(self, i):
                i += 1
                s = 0
                while i:
                    s += self.tree[i]
                    i -= i & -i
                return s
        tr = Tree(l)  # 初始化树状数组
        ans = 0
        for k in nums:
            t = k + diff
            ind = bisect_right(l, t)  # 找到 k+diff 的右边界索引
            ans += tr.presum(ind - 1)  # 查询 [0, ind-1] 范围内的元素个数
            tr.update(cnt[k])  # 更新树状数组
        return ans

Explore

是的，这种处理方式适用于所有的输入数组，包括包含负数或零的情况。计算差值的操作本质上是对每对元素进行相减，不论这些元素是正数、负数还是零，相减的操作总是有效并能正确表达两个数之间的差。因此，计算差值的步骤对于任何形式的整数输入都是有效的。

选择使用树状数组（BIT）主要是因为其在处理区间求和和元素更新的效率较高，具体来说，树状数组能在O(log n)的时间复杂度内完成这两种操作。与之相比，线段树虽然也能在O(log n)时间内完成相同的操作，但其实现更复杂，空间使用也略高。普通数组虽然更新操作只需O(1)，但计算区间和的时间复杂度为O(n)，在需要频繁执行这两种操作的情况下效率较低。因此，树状数组是一种空间和时间效率较为平衡的选择。

使用bisect_right函数查找k+diff的右边界索引时，如果k+diff恰好等于l中的某个元素，bisect_right将返回第一个大于k+diff的元素的索引。这意味着计算的是严格小于k+diff的元素的数量。因此，如果k+diff等于l中的某个元素，这个元素不会被包括在内，从而保证了数对的第二个元素严格大于第一个元素加上diff。

树状数组的update操作通过逐步增加索引i来完成，其中i在每次循环中通过i += i & -i语句递增。这个操作利用了位运算来快速找到需要更新的下一个索引。树状数组的设计确保了在初始化时数组大小足够大（通常为原数组大小加一），以容纳最大索引。此外，update操作中的循环条件i < len(self.tree)确保了索引i不会超过数组的实际边界，从而避免了越界的问题。