[315][困难][线段树][二分] 计算右侧小于当前元素的个数

题目描述

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

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

示例：

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

解题思路

二分查找

二分查找是比较直接的思路. 从右向左遍历, 将遍历过的数字存放在另外一个有序数组中, 这样对于左侧新的一个数字, 使用二分查找方法找到其对应位置的索引, 则索引值就是其有序数组中左侧元素的数量, 也就是该数字在原数组中, 对应的右侧所有小于这个数元素的数量, 将其记录在结果序列中, 并插入到有序数组中, 保持有序数组的有序性.

代码如下:

class Solution:
    def countSmaller(self, nums: List[int]) -> List[int]:
        n = len(nums)
        res, sorted_nums = [], []

        for i in range(n - 1, -1, -1):
            num = nums[i]
            index = bisect.bisect_left(sorted_nums, num)
            res.append(len(sorted_nums[:index]))
            sorted_nums.insert(index, num)

        return res[::-1]

Python中的bisect包就是二分查找的对应包.

树状数组

维护一个树状数组, 将右侧小于nums[i]元素的数量, 转换为求前缀和. 将原始数组去重后, 从小到大排列, 作为树状数组依赖的原数组, 每个数都能在这个树状数组中找到对应索引$i$, $A[i]$代表这个数出现的总次数, 对应的$C[i]$代表的就是小于等于这个数的总数量. 从后向前遍历nums, 找个当前数字对应的在树状数组中的索引$i$, 不断地更新树状数组$C$, 并从$C[i-1]$开始, 计算$0,\cdots,i-1$位置次数之和, 即为右侧小于当前数字元素的数量.

树状数组原理参考: 树状数组原理.

代码如下:

class Solution:
    def countSmaller(self, nums: List[int]) -> List[int]:
        A = list(set(nums))
        A.sort()
        m = len(A)
        C = [0] * m
        mapping = {n: i + 1 for i, n in enumerate(A)}
        res = []

        for i in range(len(nums) - 1, -1, -1):
            num = nums[i]
            c_index = mapping[num]
            while c_index <= m:
                C[c_index - 1] += 1
                c_index += c_index & (-c_index)

            t_res = 0
            c_index = mapping[num] - 1
            while c_index > 0:
                t_res += C[c_index - 1]
                c_index -= c_index & (-c_index)
            res.append(t_res)

        return res[::-1]

相关题目

• [剑指Offer-51][困难][线段树][二分] 数组中的逆序对