数据结构之树状数组浅析

刚上大一那会儿，在《算法笔记》上就读到过树状数组的内容。当时给我的感觉就是非常巧妙，还能这样设计～但平时刷题、写项目的过程中说实话用到的都不多，过一段时间就会忘记，所以索性就根据《算法笔记》上的内容，自己做一个总结，以备不时之需。

背景知识

题目

我们先来看一个非常经典的题目：

给出一个整数序列A (size = N)，接下来查询K次，每次查询给定一个正整数x (x <= N)，求前x个数之和。

这个题目不难，直接使用前缀和即可。下面我们来升级一下：假设在查询期间随时可能会对第x个数加上一个整数y，要求实时实现更新与查询操作。如果还是按照之前的做法，在每次更新的时间复杂度就为O(N)。而如果不使用前缀和的话，虽然更新的时间复杂度为O(1)，但查询的时间复杂度又变成了O(N)。无论采用哪种方法，似乎都不是非常理想。这个时候，就可以使用树状数组。

lowbit运算

在介绍树状数组之前，我们再介绍一下lowbit运算：lowbit(x) = (x) & (-x)。这个运算的实际作用就是取x二进制的最低位1及其右边的所有0。

推导过程很简单：设x的二进制有m位，其中最低位的1为第n位(0<=m<=n)。则-x根据补码的定义（每位取反后得到的数+1），对于-x的二进制第i位，若i < m，则i位依旧为0（小于m的位说明原来全部为0，取反后为1，+1后因为进位全部为又变为0）；若i == m，则i位为1（取反后为0，因为低位的进位，i位变为1，且进位不再影响高位）；若i > m，则与x相反。因此，(x) & (-x)的结果使得高于或低于n的位均为0，余下n位为1的二进制数。

树状数组

了解了lowbit运算后，我们正式介绍树状数组。对于一个数组nums，树状数组C中的每个元素存储的也是nums中某些元素之和，但不同于前缀和，树状数组第i（从1开始）个元素存储的为nums中第i个元素向前lowbit(i)个元素（包括nums[i]）之和。这么说可能比较拗口，举个例子假设nums = {1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16}，那么，对于i == 4的位置，C[4]存储的为从nums[4]开始，向前数lowbit(4) = 4个元素之和，即4+3+2+1=10。同理，C[5]存储的为从nums[5]开始，向前数lowbit(5) = 1个元素之和，即5。为了理解得更为直观，可以直接参照下图：

树状数组示意图

图中箭头表示起点，向前数对应的lowbit个数（包括起点本身）即为C中存储的元素。例如，C[10]的起点为nums[10]，向前数lowbit(10) = 2个，nums[10], nums[9]之和即为C[10]的值。注意，本篇文章为了叙述方便，树状数组的索引从1开始，实际代码中需要注意一下。

那么，树状数组究竟有什么用呢？我们还是针对之前提出的问题来一一讨论：

1. 查询函数GetSum(x)：获得nums中前x个元素之和。
2. 更新函数Update(x, val)：将第x个元素增加val。

首先来看第一个问题。假如我们要查询前x个元素，其值为sum(x) = nums[1] + nums[2] + ... + nums[x]。根据树状数组的定义，C[x] = nums[x - lowbit(x) + 1] + ... + nums[x]，可以得到：

sum(x) = nums[1] + ... + nums[x] 
	= nums[1] + ... + nums[x - lowbit(x)] +  nums[x - lowbit(x) + 1] + ... + nums[x]
	= nums[1] + ... + nums[x - lowbit(x)] + C[x]
	= sum(x - lowbit(x)) + C[x]

因此，我们可以写出如下的GetSum(x)函数：

int GetSum(int x) {
	int sum = 0;
	for (int i = x; i > 0; i -= lowbit(i) {
		sum += C[i];
	}
	return sum;
}

我们可以根据下图中的路径，以GetSum(14)（红）、GetSum(7)（黄）为例体会一下：

GetSum示意图

显然，对于GetSum(x)函数而言，求和路径为一个不断向左上行进的过程，时间复杂度O(logN)。此外，如果要求区间[x, y]内的nums元素之和，使用GetSum(y) - GetSum(x-1)即可。

第二个问题，Update(x, val)。以Update(2, val)为例，我们看下面这幅图：

Update示意图

显然，对于要更新的x，任何包含了nums[x]的C[i]我们都要更新，如图中红线所至。第一个要更新的自然为C[x]。那么，我们如何来确定下一个需要更新的C[i]呢？

从图中可以看出，下一个需要更新的C[i]为距离当前C[x]最近的能够覆盖C[x]的矩形。也就是说，lowbit[i]要至少大于lowbit[x]以覆盖C[x]。同时，我们需要C[i]尽量小，以满足距离最近的条件。所以，我们的目标就是求一个最小的y，使得

lowbit(x+y)>lowbit(x)

我们回到lowbit的定义，要想上式成立，那么x+y的二进制最低位，即最右边的1一定要在x的最右边的1左边，所以y的最低位1不能在x右边，否则x+y的二进制最低位1就会比x更小。因此，有lowbit(y) >= lowbit(x)。要使得y最小，在满足lowbit(y) >= lowbit(x)的条件下，我们让y最低位1尽量低，并且令最低位1左边的所有位均为0，此时就有y == lowbit(x)。

因此，我们可以写Update(x, val)的代码：

void Update(int x, int val) {
	for (int i = x; i <= N; i += lowbit(i)) {
		C[i] += val;
	}
}

以Update(2, val)、Update(9, val)为例，可以结合下图中的路径体会一下：
可以看到，Update(x, val)函数的时间复杂度也为O(logN)。其更新路线实际上为一个不断向右上行进的过程。

以上就是树状数组的两个核心函数，光看代码的话，还是比较清晰的。

实际运用

上面的问题只是一个引子，我们来看一个稍微复杂一点的问题：

给定一个有N个正整数的数组nums，对数组中的每个数，求出数组中该数左边比其小的数的个数，将结果作为数组返回。（0 < nums[i] < 10^5）
eg: nums{2, 6, 3, 4, 9}
return: ans{0, 1, 1, 2, 4}

这道题我们就可以使用树状数组来解决。具体方法为：由左向右遍历nums的元素x，每遍历一个x就调用update(x, 1)更新x出现的次数，同时使用GetSum(x-1)获得比x小的数的数量，具体代码如下:

#define lowbit(x) (x)&(-x) // 注意括号

class Solution {
private:
	const max_size_c_ = 100010;
	void Update(vector<int> &C, int x, int val) {
		for (int i = x; i <= N; i += lowbit(i)) {
			C[i] += val;
		}
	}
	int GetSum(vector<int> &C, int x) {
		int sum = 0;
		for (int i = x; i > 0; i -= lowbit(i) {
			sum += C[i];
		}
		return sum;
	}
public:
    vector<int> GetNumberLessFromLeft(const vector<int>& nums) {
        int n = nums.size();
        vector<int> ans(n, 0), C(max_size_c_, 0)
        for (int i = 0; i < n; i++) {
			update(C, nums[i], 1);
			ans[i] = GetSum(C, nums[i] - 1);
		}
        return ans;
    }
};

如果要求左边比x大的数，使用GetSum(C, max_size_c_) - GetSum(C, nums[i] - 1)即可。

这里算法笔记写上的是GetSum(C, nums.size()) - GetSum(C, nums[i] - 1)，我感觉应该用max_size_c_，因为统计比nums[i]大，那么应该令最大的数作为边界。欢迎讨论～

结语

可以看到，树状数组的定义非常巧妙，能够将O(N)的时间复杂度降低至O(logN)。也因此，我非常喜欢树状数组这一数据结构。但自大一以来，其实树状数组的使用并不是很多。有时候如果为了用而用，没有对实际需求/问题起到实质性的帮助，反而有点舍本逐末的意思。归根结底，树状数组也只是一个数据结构，一个工具。我们有时候并不应该拿着锤子找钉子，而是时刻备好这把备受偏爱的锤子，当特定钉子出现的时候能够用得上。

后续如果有时间会更新其他定义的树状数组。

2023-12-3 22:42

参考

《算法笔记》胡凡、曾磊