树状数组-归并排序-逆序对-2426. 满足不等式的数对数目

问题描述

给你两个下标从 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 <= 105
-104 <= nums1[i], nums2[i] <= 104
-104 <= diff <= 104

问题求解

首先对问题做一下转化，等价于求解nums[i] <= nums[j] + diff的个数。

• 归并排序
  归并排序是求解逆序对问题的经典解法，优点是容易思考，编码也相对容易，且对数据大小没有限制，需要优先掌握。

class Solution:
    # nums[i] <= nums[j] + diff
    def numberOfPairs(self, nums1: List[int], nums2: List[int], diff: int) -> int:
        n = len(nums1)
        nums = [nums1[i] - nums2[i] for i in range(n)]

        def merge_sort(nums):
            if len(nums) <= 1:
                return 0
            res = 0
            mid = len(nums) // 2
            l, r = nums[:mid], nums[mid:]
            res += merge_sort(l)
            res += merge_sort(r)

            i, n, m = 0, len(l), len(r)
            for x in r:
                while i < n and l[i] <= x + diff: i += 1
                res += i
            
            i, j, k = 0, 0, 0
            while i < n and j < m:
                if l[i] <= r[j]:
                    nums[k] = l[i]
                    i += 1
                else:
                    nums[k] = r[j]
                    j += 1
                k += 1
            while i < n:
                nums[k] = l[i]
                i += 1
                k += 1
            while j < m:
                nums[k] = r[j]
                j += 1
                k += 1
            return res
        
        return merge_sort(nums)

• 树状数组
  转化为计数问题。

from sortedcontainers import SortedSet
class Solution:
    # nums[i] <= nums[j] + diff
    def numberOfPairs(self, nums1: List[int], nums2: List[int], diff: int) -> int:
        nums = [x - y for x, y in zip(nums1, nums2)]
        
        nums_sort = SortedSet()
        for num in nums:
            nums_sort.add(num)
            nums_sort.add(num + diff)
        record = {}
        rank = 1
        for num in nums_sort:
            record[num] = rank
            rank += 1

        bit = [0] * (len(record) + 1)

        def query(i):
            res = 0
            while i:
                res += bit[i]
                i -= i & -i
            return res
        
        def update(i):
            while i < len(bit):
                bit[i] += 1
                i += i & -i
        
        res = 0
        for num in nums:
            res += query(record[num + diff])
            update(record[num])
        
        return res