树状数组-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