树状数组-315. 计算右侧小于当前元素的个数

问题描述

给你一个整数数组 nums ，按要求返回一个新数组 counts 。数组 counts 有该性质： counts[i] 的值是  nums[i] 右侧小于 nums[i] 的元素的数量。

示例 1：

输入：nums = [5,2,6,1]
输出：[2,1,1,0]
解释：
5 的右侧有 2 个更小的元素 (2 和 1)
2 的右侧仅有 1 个更小的元素 (1)
6 的右侧有 1 个更小的元素 (1)
1 的右侧有 0 个更小的元素
示例 2：

输入：nums = [-1]
输出：[0]
示例 3：

输入：nums = [-1,-1]
输出：[0,0]

提示：

1 <= nums.length <= 105
-104 <= nums[i] <= 104

问题抽象

求解满足nums[j] < nums[i] and j > i的j的个数。
实际上对于求解 a1 * nums[j] + b1 < a2 * nums[i] * n + b2 and j > i的j的个数都可以用排序离散化、计频数组、区间求和来解决。

问题求解

• 离散化：由于需要比较大小，但具体数值其实并不重要，因此可以对原数据进行离散化处理，降低桶的空间复杂度；
• 区间求和：从末尾开始遍历并对桶做更新，同时计数比当前数小的桶里数的总数，这个可以用树状数组高效求解；

class Solution:
    def countSmaller(self, nums: List[int]) -> List[int]:
        n = len(nums)
        rank = 1
        record = {}
        sort = sorted(nums)
        for idx in range(n):
            if idx and sort[idx] == sort[idx - 1]:
                continue
            record[sort[idx]] = rank
            rank += 1
        bit = [0] * (len(record) + 1)

        def update(idx, delta):
            while idx < len(bit):
                bit[idx] += delta
                idx += (idx & -idx)
        
        def query(idx):
            res = 0
            while idx:
                res += bit[idx]
                idx -= (idx & -idx)
            return res
        
        res = [0] * n
        for idx in range(n - 1, -1, -1):
            rank = record[nums[idx]]
            update(rank, 1)
            res[idx] = query(rank - 1)
        return res