什么是树状数组
我们知道,前缀和数组能解决任意一段区间的累加和问题
但这建立在数组中的元素不发生变化的情况,如果可以修改原始数组中的某个元素,为了让前缀和数组正确,就需要在前缀和数组中修改该元素位置后面的所有的数,时间复杂度为O(N)
而树状数组能做到查询区间和,修改单个元素都为O(logN)
前缀和 | 树状数组 | |
---|---|---|
区间查询 | O(1) | O(logN) |
修改单个元素 | O(N) | O(logN) |
因此,树状数组专门解决带单点更新的区间累加和需求
结构
对于长度为n的原数组data,生成一个长度为n+1的tree数组
public class IndexTree {
// 原始数组
private int[] data;
// tree数组
private int[] tree;
private int size;
public IndexTree(int[] data) {
this.data = data;
this.size = data.length;
this.tree = new int[size+1];
}
}
为什么tree长度为n+1?第tree数组中第0位不用
tree中的每一项i,其实代表一个范围的累加和,代表哪个范围呢?
将i的二进制数,抹去最后一个1,再加1,记为newi
代表原始数组中第newi个数,到第i个数和
什么叫抹去最后一个1?即减去二进制下从右往左数的第一个1
例如1111抹去最后一个1变为1110
10110抹去最后一个1变为10100
例如:i = 12
- 其二进制表示为1100
- 抹去最后一个1(即100)为1000
- 加1得到1001
那么tree[i]的值为,原数组中第1001(十进制为9)个数到第i个数,也就是1100(十进制为12),这些数的和
根据这个规则,我们看看从tree中,下标为1到16的数,代表原数组中哪些数的累加和
tree中的下标i | 下标的二进制表示 | 减去最后一个1 | 再加1 | 下标本身 | 含义(二进制表示) |
---|---|---|---|---|---|
1 | 1 | 0 | 1 | 1 | 第1个数到第1个数的和 |
2 | 10 | 0 | 1 | 10 | 第1个数到第10个数的和 |
3 | 11 | 10 | 11 | 11 | 第11个数到第11个数的和 |
4 | 100 | 0 | 1 | 100 | 第1个数到第100个数的和 |
5 | 101 | 100 | 101 | 101 | 第101个数到第101个数的和 |
6 | 110 | 100 | 101 | 110 | 第101个数到第110个数的和 |
7 | 111 | 110 | 111 | 111 | 第111个数到第111个数的和 |
8 | 1000 | 0 | 1 | 1000 | 第1个数到第1000个数的和 |
9 | 1001 | 1000 | 1001 | 1001 | 第1001个数到第1001个数的和 |
10 | 1010 | 1000 | 1001 | 1010 | 第1001个数到第1010个数的和 |
11 | 1011 | 1010 | 1011 | 1011 | 第1011个数到第1011个数的和 |
12 | 1100 | 1000 | 1001 | 1100 | 第1001个数到第1100个数的和 |
13 | 1101 | 1100 | 1101 | 1101 | 第1101个数到第1101个数的和 |
14 | 1110 | 1100 | 1101 | 1110 | 第1101个数到第1110个数的和 |
15 | 1111 | 1110 | 1111 | 1111 | 第1111个数到第1111个数的和 |
16 | 10000 | 0 | 1 | 10000 | 第1个数到第10000个数的和 |
前缀和
如果我想求原始数组中,从第一个数开始到第i个数的累加和,应该怎么求?
假设i为45,其二进制表示为101101
准备一个累加和变量res
- 首先从tree数组中找到下标为101101的值,加到res中,即res += sum[101101]
- 将
i
抹去最后一个1,变为101100,res += tree[101100]
- 再抹去最后一个1,变为101000,
res += tree[101000]
- 再抹去最后一个1,变为100000,
res += tree[100000]
此时如果再抹去最后一个1,i将变为0,所以停止
总结规律:不断抹去i的最后一个1,sum[i]累加到结果中
public int sum(int i) {
int res = 0;
while (i > 0) {
res += tree[index];
// 抹去最右侧的1
i -= i & (~i + 1);
}
return res;
}
正确性证明
为什么这么做,能正确计算出前缀和呢?
以i = 45 (101101)为例,我们依次看抹去最后一个1后的值在sum数组中代表什么:
i | sum[i]表示的起始位置 | sum[i]表示的结束位置 | |
---|---|---|---|
没有抹去 | 101101 | 101101 | 101101 |
第一次抹去 | 101100 | 101001 | 101100 |
第二次抹去 | 101000 | 100001 | 101000 |
第三次抹去 | 100000 | 1 | 100000 |
可以发现,这4次的i值在sum数值中所代表的的区间和,不重不漏地覆盖了从1到101101的所有数:
- sum[100000]:从第1个数到第100000个数的累加和
- sum[101000]:从第100001个数到第101000个数的累加和
- sum[101100]:从第101001个数到第101100个数的累加和
- sum[101101]:从第101101个数到第101101个数的累加和
将sum数组中这4个数累加起来,恰好就能得到从第1到第101101个数的前缀和
时间复杂度
每次while循环抹去最右侧的1,最多抹去logN次,因此时间复杂度为O(logN)
单点增加值
假设将i位置的数加上v
当修改原始数组中某个数时,需要同时修改sum数组,怎么知道在sum数组中哪些数受牵连呢?
例如当size = 16,我修改第3个数时,3的二进制表示为11
- 将11加上最右侧的1,得到100,sum[110] += v
- 将100加上最右侧的1,得到1000,sum[1000] += v
- 将1000加上最右侧的1,得到10000,大于size,结束循环
这些位置的数,都是因为原始数组中第i个数变化了,需要调整的位置
总结规律:不断将i加上最后一个1,sum[i] += v,直到i大于size为止
public void add(int i,int v) {
while (i <= size) {
tree[i] += v;
// i加上最右侧的1
i += i & (~i + 1);
}
}
时间复杂度
每次while循环加上最右侧的1,其实没加几次就会开始每次循环翻倍,最多加logN 次,因此时间复杂度为O(logN)
初始化tree数组
根据原始数组初始化tree数组时,复用add方法就行:
先假设原始数组全为0,依次给每个位置i
增加data[i]
的值,就等于初始化好了tree数组
public IndexTree(int[] data) {
this.data = data;
this.size = data.length;
this.tree = new int[size+1];
for (int i = 1;i<=size;i++) {
add(i, data[i-1]);
}
}
单点修改值
现在可以实现将某个点的值增加v
,也可以复用该方法将某个点的值修改成d
- 先计算d和原始值的
差
- 调用add方法将原始值增加这个差
public void set(int i,int d) {
int diff = d - data[i-1];
add(i, diff);
}
计算区间和
有了前缀和,计算区间和的就方便了
public int rangeSum(int left,int right) {
return sum(right) - sum(left-1);
}