「树状数组」第 4 节： 相关例题

其实下面这两个问题本质上是一个问题。

例 1：《剑指 Offer 》第 51 题：逆序数的计算

• 剑指 Offer 51. 数组中的逆序对。

在数组中的两个数字如果前面一个数字大于后面的数字，则这两个数字组成一个逆序对。

输入一个数组，求出这个数组中的逆序对的总数。

样例

输入：[1,2,3,4,5,6,0]

输出：6

分析：这道题最经典的思路是使用分治法计算，不过使用树状数组语义更清晰一些。

Python 代码：

class Solution(object):
    def inversePairs(self, nums):
        class FenwickTree:
            def __init__(self, n):
                self.size = n
                self.tree = [0 for _ in range(n + 1)]

            def __lowbit(self, index):
                return index & (-index)

            # 单点更新：从下到上，最多到 size，可以取等
            def update(self, index, delta):
                while index <= self.size:
                    self.tree[index] += delta
                    index += self.__lowbit(index)

            # 区间查询：从上到下，最少到 1，可以取等
            def query(self, index):
                res = 0
                while index > 0:
                    res += self.tree[index]
                    index -= self.__lowbit(index)
                return res

        # 特判
        l = len(nums)
        if l < 2:
            return 0

        # 原始数组去除重复以后从小到大排序
        s = list(set(nums))

        # 构建最小堆，因为从小到大一个一个拿出来，用堆比较合适
        import heapq
        heapq.heapify(s)

        # 由数字查排名
        rank_map = dict()
        index = 1
        # 不重复数字的个数
        size = len(s)
        for _ in range(size):
            num = heapq.heappop(s)
            rank_map[num] = index
            index += 1

        res = 0
        # 树状数组只要不重复数字个数这么多空间就够了
        ft = FenwickTree(size)
        # 从后向前看，拿出一个数字来，就更新一下，然后向前查询比它小的个数
        for i in range(l - 1, -1, -1):
            rank = rank_map[nums[i]]
            ft.update(rank, 1)
            res += ft.query(rank - 1)
        return res

说明：中间将数字映射到排名是将原数组「离散化」，「离散化」的原因有 2 点：

1、树状数组我们看到，下标是从 $1$ 开始的，我们不能保证我们的数组所有的元素都大于等于 $1$；

2、即使元素都大于等于 $1$，为了节约树状数组的空间，我们将之「离散化」可以把原始的数都压缩到一个小区间。我说的有点不太清楚，这一点可以参考 树状数组 求逆序数 poj 2299。

例2：「力扣」第 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 个更小的元素.

分析：事实上，这个问题就是在计算「逆序数」，和上一个问题是一样的。「计算右侧小于当前元素的个数」我们可以「从后向前一个一个填」。因为涉及大小关系，所以要排个序，并且给出排名（rank）。这一步操作也叫「离散化」。具体方法是：先画出一个排名表，对于这个问题，排名表是：

数	排名
$5$	$3$
$2$	$2$
$1$	$1$
$6$	$4$

从后向前填：

1、遇到 $1$ ，$1$ 的排名是 $1$ ，首先先在 $1$ 那个位置更新 $1$，那么 $1$ 之前肯定没有数了，所以就是 $0$；

2、遇到 $6$ ， $6$ 的排名是 $4$，首先先在 $4$ 那个位置更新 $1$，那么 $6$ 之前可以在树状树组里面查一下，是 $1$；

3、遇到 $2$ ， $2$ 的排名是 $2$，首先先在 $2$ 那个位置更新 $1$，那么 $2$ 之前可以在树状树组里面查一下，是 $1$；

4、遇到 $5$ ，$5$ 的排名是 $3$，首先先在 $3$ 那个位置更新 $1$，那么 $3$ 之前可以在树状树组里面查一下，是 $2$；

反过来就是结果 $[2,1,1,0]$。

Python 代码：

class Solution:
    def countSmaller(self, nums):
        """
        :type nums: List[int]
        :rtype: List[int]
        """

        class FenwickTree:
            def __init__(self, n):
                self.size = n
                self.tree = [0 for _ in range(n + 1)]

            def __lowbit(self, index):
                return index & (-index)

            # 单点更新：从下到上，最多到 size，可以取等
            def update(self, index, delta):
                while index <= self.size:
                    self.tree[index] += delta
                    index += self.__lowbit(index)

            # 区间查询：从上到下，最少到 1，可以取等
            def query(self, index):
                res = 0
                while index > 0:
                    res += self.tree[index]
                    index -= self.__lowbit(index)
                return res

        l = len(nums)
        if l == 0:
            return []
        if l == 1:
            return [0]

        s = list(set(nums))
        import heapq
        heapq.heapify(s)
        index = 1

        size = len(s)
        rank_map = dict()
        ft = FenwickTree(size)
        for _ in range(size):
            num = heapq.heappop(s)
            rank_map[num] = index
            index += 1

        # 从后向前填表
        res = [None for _ in range(l)]
        for index in range(l - 1, -1, -1):
            rank = rank_map[nums[index]]
            ft.update(rank, 1)
            res[index] = ft.query(rank - 1)
        return res

（本文完）