315. 计算右侧小于当前元素的个数

Difficulty
Hard
Tags
树状数组
URL
https://leetcode-cn.com/problems/count-of-smaller-numbers-after-self/
Star
给你一个整数数组 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

树状数组

思路
将该问题抽象为树状数组题目。
额外维护一个数组tree ，数组值为 该下标表示的值出现次数。
然后从后向前遍历数组nums，遍历到元素num时候，
  • 先计算 右侧小于num 的个数，即 计算一下 数组tree[:num] 范围和
  • 然后更新数组tree，即tree[num] += 1
至此，便可将问题抽象为 树状数组问题。
技巧：对数组nums 进行一个映射，因为我其实并不关心nums中的具体值，我之关心元素相对大小。所以可以对nums进行一个压缩，示例如下：
notion imagenotion image
为什么要压缩呢？因为我们要把nums中的值作为 额外数组 tree 的下标，所以压缩之后，节省空间。比如：nums = [100], 如果不压缩，那么 tree 长度也是100，但是压缩之后，只需要1就可以。
题解
class Solution:
    def countSmaller(self, nums: List[int]) -> List[int]:
        
        # 数组压缩，压缩的时候可以去掉重复元素
        uniques = sorted(set(nums))
        hash_map = {num: i for i, num in enumerate(uniques)}
        
        res = [0 for _ in nums]
        t = TreeNum(len(uniques))
        for i in range(len(nums)-1, -1, -1):
            # 获取压缩后的值
            x = hash_map[nums[i]]
            # 1. 取得当前小于 x 的个数
            res[i] = t.prefix(x-1)
            # 2. 更新树状数组
            t.update(x, 1)
        
        return res

class TreeNum:
    def __init__(self, n):
        self.tree = [0] * (n+1) # 前面多放一个 0 

    def update(self, idx, delta):
        idx += 1
        while idx < len(self.tree):
            self.tree[idx] += delta
            idx += self.lowbit(idx)

    def prefix(self, idx):
        idx += 1
        res = 0
        while idx > 0:
            res += self.tree[idx]
            idx -= self.lowbit(idx)

        return res

    def lowbit(self, x):
        return x & (-x)