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

标签: 树状数组 线段树 数组 二分查找 分治 有序集合 归并排序

难度: Hard

给你一个整数数组 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

Submission

运行时间: 561 ms

内存: 35.1 MB

class Solution:
    def countSmaller(self, nums: List[int]) -> List[int]:

        bucket ={n:i for i, n in enumerate(sorted(set(nums)), start=1)}
        n = len(bucket) + 1
        tree = [0]*n


        def prefix_sum(i):
            ans = 0
            while i > 0:
                ans += tree[i]
                i = i & (i-1)
            return ans
        def add(i):
            while i < n:
                tree[i] += 1
                i += (i & -i)
            
        ans = []

        for x in nums[::-1]:
            i = bucket[x]
            ans.append(prefix_sum(i-1))
            add(i)
        return ans[::-1]





        

Explain

该题解使用树状数组(Binary Indexed Tree)来解决问题。首先将原数组去重并排序,建立数值到索引的映射。然后从右往左遍历原数组,对于每个元素,查询比它小的元素个数即为树状数组中对应索引前面的前缀和,然后将当前元素对应的索引位置加1,表示当前元素已被计入答案。最后将答案数组反转输出。

时间复杂度: O(nlogn)

空间复杂度: O(n)

class Solution:
    def countSmaller(self, nums: List[int]) -> List[int]:
        
        # 将去重后的元素映射到从1开始的索引
        bucket ={n:i for i, n in enumerate(sorted(set(nums)), start=1)}
        n = len(bucket) + 1
        
        # 初始化大小为n的树状数组
        tree = [0]*n

        # 查询前缀和
        def prefix_sum(i):
            ans = 0
            while i > 0:
                ans += tree[i]
                i = i & (i-1)
            return ans
        
        # 更新树状数组
        def add(i):
            while i < n:
                tree[i] += 1
                i += (i & -i)
            
        ans = []

        # 从右往左遍历数组
        for x in nums[::-1]:
            # 查询比当前元素小的个数
            i = bucket[x]
            ans.append(prefix_sum(i-1))
            # 将当前元素加入树状数组 
            add(i)
        
        # 将结果反转后返回
        return ans[::-1]

Explore

在题解中,数组中的重复元素通过使用`set`去重,然后再对结果排序的方式处理。这样做的目的是将每个数值映射到一个唯一的索引,这样可以在树状数组中正确地跟踪每个数值的出现频率。去重是必要的,因为我们需要对每个唯一的数值在树状数组中维护一个计数,而排序则是为了确保树状数组中的索引可以代表小于当前数值的所有数值的累计频率。

树状数组中的`prefix_sum(i)`函数通过累加树状数组中的值来计算索引`i`之前的所有元素的总和。在函数中,利用了树状数组的性质,通过循环中的`i = i & (i-1)`操作,逐步移除最低位的1,从而向前移动到下一个要累加的区间。这种方式确保了只计算到索引`i`之前的元素总和,因为在树状数组中,每个索引位置包含了从其能覆盖的最右侧区间起到自己的累计和。

在题解中,树状数组的大小`n`是通过计算去重并排序后的数组的长度得出的,这确保了每个唯一数值都有一个对应的索引。加1的原因是树状数组的索引从1开始,因此为了容纳所有索引,数组的长度需要是去重后的数值个数加上1。这样,数组的索引0不会被使用,从而避免了索引混淆并简化了树状数组的操作。