深入理解数据结构 —— 树状数组

什么是树状数组

我们知道，前缀和数组能解决任意一段区间的累加和问题

但这建立在数组中的元素不发生变化的情况，如果可以修改原始数组中的某个元素，为了让前缀和数组正确，就需要在前缀和数组中修改该元素位置后面的所有的数，时间复杂度为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

1. 先计算d和原始值的差
2. 调用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);
}