美团240316春招实习笔试——小美的逆序对（树状数组）

题目描述

小美拿到了一个排列，她定义$f(i)$为:将第$i$个元素取反后，形成的数组的逆序对数量。

逆序对：对于一个整数数组 nums，逆序对是一对满足 0 <= i < j < nums.length 且 nums[i] > nums[j]的整数对 [i, j] 。
排列：排列是指一个长度为n的数组，1到n每个元素恰好出现了一次。

小美希望你求出$f(1)$到$f(n)$的值。

输入描述

第一行输入一个正整数n，代表排列的大小
第二行输入n个正整数$a_i$，代表排列的元素。

约束

$$\begin{gathered} 1\le n\le 2\times 10^5 \\ 1\le a_i\le n \end{gathered}$$

输出描述

输出$n$个整数，第$i$个整数是$f(i)$的值。

样例

输入

3
1 2 3

输出

0 1 2

说明

第一个元素取反，数组将变成[-1,2,3]，逆序对数量为 0.
第二个元素取反，数组将变成[1,-2,3]，逆序对数量为1.
第三个元素取反，数组将变成[1,2,-3]，逆序对数量为2.

解题

定义两个数组leftInverseNums和rightInverseNums，分别记录nums中下标为 $i$ 的元素 $num$ 与其左侧和右侧的数字构成的逆序对数量。
数组nums原本的逆序对数量可由两个数组中的任一获得：totalInverseNum = sum(leftInverseNums)。
假设将下标为 $i$ 的元素 $num$ 取反。

• 对于 $num$ 左侧的元素，$-num$ 会它们都更小，因此左侧均构成逆序对：totalInverseNum += i；
  • 因为totalInverseNum已经统计过左侧元素了，避免重复统计，要减去左侧原有的逆序对数：totalInverseNum -= leftInverseNums;
• 对于 $num$ 右侧的元素，原本比 $num$ 小的值不再比 $-num$小，右侧的逆序对不再成立：rightInverseNums -= rightInverseNums。

而计算leftInverseNums和rightInverseNums的过程可以使用树状数组进行：

bitLeft = BIT(n)  # 根据题设，数字分布于[1,n]区间，bit[i]表示区间[1,i]有多少数字已被遍历
leftInverseNums = [0 for _ in range(n)]  # 记录nums[i]和左侧数字构成的逆序对数
for i, num in enumerate(nums):
	# [num+1, n] = [1, n] - [1, num]，表示此前有多少大于num的数字，即num和它之前的数字构成的逆序对数量。
    val = bitLeft.query(n) - bitLeft.query(num)  
    leftInverseNums[i] = val
    bitLeft.add(num, 1)  # 关键：遍历过的位置置1，那么区间和就是区间内遍历过的元素个数

右侧也是如此：

bitRight = BIT(n)
rightInverseNums = [0 for _ in range(n)]  # 记录nums[i]和右侧数字构成的逆序对数
# 从右向左查询比当前num小的数字个数，即num右侧的逆序对数量。这些逆序对在num取反后不再逆序。
for i, num in reversed(list(enumerate(nums))):  
    val = bitRight.query(num-1)  # 比num小的数字即属于区间[1, num-1]
    rightInverseNums[i] = val
    bitRight.add(num, 1)

完整代码如下：

n = 4
nums = [1, 4, 3, 2]

# 树状数组
class BIT:
    def __init__(self, n):
        self.n = n
        self.tree = [0 for _ in range(n+1)]  # 下标从1开始

    def lowBit(self, x: int) -> int:
        return x & (-x)

    def add(self, i: int, val: int):
        while i <= self.n:
            self.tree[i] += val
            i += self.lowBit(i)  # i的父节点下标为i+lowbit(i)
        
    def query(self, i: int = 1) -> int:  # 查询（0，i]区间内的所有元素的和
        res = 0
        while i > 0:
            res += self.tree[i]
            i -= self.lowBit(i)  # 下一个要查找的坐标为i-lowbit(i)
        return res

bitLeft = BIT(n)  # 根据题设，数字分布于[1,n]区间，bit[i]表示区间[1,i]有多少数字已被遍历
leftInverseNums = [0 for _ in range(n)]  # 记录nums[i]和左侧数字构成的逆序对数
for i, num in enumerate(nums):
    val = bitLeft.query(n) - bitLeft.query(num)  # [num+1, n] = [1, n] - [1, num]，表示此前有多少大于num的数字，即num和它之前的数字构成的逆序对数量。
    leftInverseNums[i] = val
    bitLeft.add(num, 1)  # 关键：遍历过的位置置1，那么区间和就是区间内遍历过的元素个数

bitRight = BIT(n)
rightInverseNums = [0 for _ in range(n)]  # 记录nums[i]和右侧数字构成的逆序对数
for i, num in reversed(list(enumerate(nums))):  # 从右向左查询比当前num小的数字个数，即num右侧的逆序对数量。这些逆序对在num取反后不再逆序。
    val = bitRight.query(num-1)  # 比num小的数字即属于区间[1, num-1]
    rightInverseNums[i] = val
    bitRight.add(num, 1)

totalInverse = sum(rightInverseNums)  
# print(bitLeft.tree)
# print(bitRight.tree)
# print(leftInverseNums)
# print(rightInverseNums)
for i in range(n):
    print(totalInverse - rightInverseNums[i] + i - leftInverseNums[i], end = ' ')

参考文章