二叉索引树（树状数组）的原理

发表于 | 更新于 | 分类于 | 阅读次数：

背景

了解到二叉索引树这个数据结构，是在 leetcode 的 307 题，题目是要求实现一个数据结构，可以返回数组任意区间的和以及更新数组的某个值。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
307、Range Sum Query - Mutable

Given an integer array nums, find the sum of the elements between indices i and j (i ≤ j), inclusive.

The update(i, val) function modifies nums by updating the element at index i to val.

Example:

    Given nums = [1, 3, 5]
    
    sumRange(0, 2) -> 9
    update(1, 2)
    sumRange(0, 2) -> 8

Constraints:

- The array is only modifiable by the update function.
- You may assume the number of calls to update and sumRange function is distributed evenly.
- 0 <= i <= j <= nums.length - 1

常规解法

先介绍下常规的解法，树状数组有用到他们之中的一些思想或者过程。

解法一

最暴力的解法，sumRange 直接 for 循环算，update 直接更新数组中的值。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
/**
 * @param {number[]} nums
 */
var NumArray = function (nums) {
  this.nums = [...nums];
};

/**
 * @param {number} i
 * @param {number} val
 * @return {void}
 */
NumArray.prototype.update = function (i, val) {
  this.nums[i] = val;
};

/**
 * @param {number} i
 * @param {number} j
 * @return {number}
 */
NumArray.prototype.sumRange = function (i, j) {
  let sum = 0;
  for (let k = i; k <= j; k++) {
    sum += this.nums[k];
  }
  return sum;
};

/**
 * Your NumArray object will be instantiated and called as such:
 * var obj = new NumArray(nums)
 * obj.update(i,val)
 * var param_2 = obj.sumRange(i,j)
 */

时间复杂度： update 是 O(1)，sumRange 是 O(n)。

解法二

303 题 做过 sumRange 的优化，我们用一个数组保存累计的和，numsAccumulate[i] 存储 0 到 i - 1 累计的和。

如果我们想求 i 累积到 j 的和，只需要用 numsAccumulate[j + 1] 减去 numsAccumulate[i]。

结合下边的图应该很好理解，我们要求的是橙色部分，相当于 B 的部分减去 A 的部分。

img

所以我们可以提前把一些前缀和存起来，然后查询区间和的时候在可以通过差实现。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
/**
 * @param {number[]} nums
 */
var NumArray = function (nums) {
  this.nums = [...nums];
  this.numsAccumulate = [0];
  let sum = 0;
  for (let i = 0; i < nums.length; i++) {
    sum += nums[i];
    this.numsAccumulate.push(sum);
  }
};

/**
 * @param {number} i
 * @param {number} val
 * @return {void}
 */
NumArray.prototype.update = function (i, val) {
  let sub = val - this.nums[i];
  this.nums[i] = val;
  for (let k = i + 1; k < this.numsAccumulate.length; k++) {
    this.numsAccumulate[k] += sub;
  }
};

/**
 * @param {number} i
 * @param {number} j
 * @return {number}
 */
NumArray.prototype.sumRange = function (i, j) {
  return this.numsAccumulate[j + 1] - this.numsAccumulate[i];
};

/**
 * Your NumArray object will be instantiated and called as such:
 * var obj = new NumArray(nums)
 * obj.update(i,val)
 * var param_2 = obj.sumRange(i,j)
 */

时间复杂度： update 是 O(n)，sumRange 是 O(1)。

虽然 sumRange 的时间复杂度优化了，但是 update 又变成了 O(n)。因为更新一个值的时候，这个值后边的累计和都需要更新。

解法三

解法一和解法二时间复杂度两个方法始终一个是 O(1)，一个是 O(n)。这里再分享 官方题解 提供的一个解法，可以优化查询区间的时间复杂度。

我们可以将原数据分成若干个组，然后提前计算这些组的和，举个例子。

1
2
3
4
组号:    0         1            2              3
数组: [2 4 5 6] [9 9 3 8] [1  2  3   4]  [4    2  3   4]
下标:  0 1 2 3   4 5 6 7   8  9  10  11   12  13  14  15
和:      17        29          10              13

如果我们要计算 sumRange(1,13)，之前我们需要循环累加下标 1 到 13 的数字的和。

现在我们只需要循环累加 1 到 3 的和，加上循环累加 12 到 13 的和，再累加中间组提前算好的和，也就是第 1 组和第 2 组的和 29 和 10 ，就是最终的结果了。

至于更新的话，我们也不需要像解法二那样更新那么多。我们只需要更新当前元素所在的组即可。

下一个问题，每组的大小定多少呢？

如果定的小了，那么组数就会特别多。

如果定的大了，那么组内元素就会特别多。

组数和组内元素个数都会影响到 sumRange 的时间复杂度。

这里，我们在组数和组内元素个数之间取个平衡，假设数组大小是 n，那么组内元素个数取 $\sqrt{n}$ ，这样的话组数也是 $\sqrt{n}$ ，这样就可以保证我们查询的时间复杂度是 $O(\sqrt{n})$ 了。因为最坏的情况，无非是查询范围跨越整个数组，中间我们需要累加 $\sqrt{n} - 2$ 个组，第 0 组最多累加 $\sqrt{n}$ 次，最后一组也最多累加 $\sqrt{n}$ 次，整体上就是 $O(\sqrt{n})$ 了。

结合代码理解一下。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
/**
 * @param {number[]} nums
 */
var NumArray = function (nums) {
  this.nums = [...nums];
  this.groupSize = Math.floor(Math.sqrt(this.nums.length));
  this.group = [];
  let sum = 0;
  let i = 0;
  for (i = 0; i < nums.length; i++) {
    sum += nums[i];
    if ((i + 1) % this.groupSize === 0) {
      this.group.push(sum);
      sum = 0;
    }
  }
  //有可能数组大小不能整除组的大小, 最后会遗漏下几个元素
  if (i % this.groupSize !== 0) {
    this.group.push(sum);
  }
};

/**
 * @param {number} i
 * @param {number} val
 * @return {void}
 */
NumArray.prototype.update = function (i, val) {
  let sub = val - this.nums[i];
  let groudId = Math.floor(i / this.groupSize);
  this.group[groudId] += sub;
  this.nums[i] = val;
};

/**
 * @param {number} i
 * @param {number} j
 * @return {number}
 */
NumArray.prototype.sumRange = function (i, j) {
  let groupI = Math.floor(i / this.groupSize);
  let groupJ = Math.floor(j / this.groupSize);
  let sum = 0;
  //在同一组内, 直接累加
  if (groupI === groupJ) {
    for (let k = i; k <= j; k++) {
      sum += this.nums[k];
    }
  } else {
    //左边组的元素累加
    for (let k = i; k < (groupI + 1) * this.groupSize; k++) {
      sum += this.nums[k];
    }
    //累加中间所有的组
    for (let g = groupI + 1; g < groupJ; g++) {
      sum += this.group[g];
    }
    //右边组的元素累加
    for (let k = groupJ * this.groupSize; k <= j; k++) {
      sum += this.nums[k];
    }
  }
  return sum;
};

/**
 * Your NumArray object will be instantiated and called as such:
 * var obj = new NumArray(nums)
 * obj.update(i,val)
 * var param_2 = obj.sumRange(i,j)
 */

时间复杂度： update 是 O(1)，sumRange 是 $O(\sqrt{n})$ 。

树状数组

有了上边的背景，我们再回到树状数组。

这个解法写法很简单，但理解的话可能稍微难一些，很多文章都直接讲该怎么用，没有介绍最初的动机，于是去看了提出这个解法的 原始论文，看看能不能理解。

这个解法叫 Fenwick tree 或者binary indexed tree，翻译过来的话叫做树状数组或者二叉索引树，但我觉得 binary 翻译成二进制更好，叫做二进制索引树更贴切些，二叉树容易引起误解。

回想一下解法三，我们预先求出了若干个区间和，然后查询的区间可以根据之前预先求出来的区间来求出。这里的话同样的思想，先预先求一些区间和，然后把要求的区间分解成若干个之前求好的区间和即可。相比于解法三，这里的分解会更加巧妙一些。

我们知道计算机中的数都是由二进制来表示的，任何一个数都可以分解成 2 的幂次的和，进制转换不熟的话可以参考 再谈进制转换。

举个例子 $11 = 2^0 + 2^1 + 2^3 = 1 + 2 + 8$，$9=2^0+2^3=1+8$ 等等。

接下来就是神奇的地方了，每一个数都可以拆成这样的 x = a + b + c + ... 的形式。

我们把等式左侧的数 x 看做是区间 [1, x]，等式右边看做从 x 开始每个区间的长度，也就变成了下边的样子。

[1, x] = [x, x - a + 1] + [x - a, x - a - b + 1] + [x - a - b, x - a - b - c + 1] + ...。

看起来有些复杂，举个具体的例子就简单多了。

以 $11 = 2^0 + 2^1 + 2^3 = 1 + 2 + 8$ 为例，可以转换为下边的等式。

[1, 11] = [11, 11] + [10, 9] + [8, 1]。

[11, 11]、[10, 9]、[8, 1] 长度分别是 1、2、8。

我们成功把一个大区间，分成了若干个小区间，这就是树状数组最核心的地方了，只要理解了上边讲的，下边就很简单了。

首先，因为数组的下标是从 0 开始的，上边的区间范围是从 1 开始的，所以我们在原数组开头补一个 0 ，这样区间就是从 1 开始了。

因此我们可以通过分解快速的求出 [1, x] 任意前缀区间的和，知道了前缀区间的和，就回到了解法二，通过做差可以算出任意区间的和了。

最后，我们需要解决子区间该怎么求？

[1, 11] = [11, 11] + [10, 9] + [8, 1] 我们用 V 表示子区间，用 F 表示某个区间。

F[1,11] = V[11] + V[10] + V[8]

其中，V[11] = F[11,11], V[10] = F[10,9], V[8]=F[8...1]，为什么是这样？

回到二进制，F[0001,1011] = V[1011] + V[1010] + V[1000]

1010 = 1011 - 0001，0001 就是十进制的 1，所以 V[1011] 存 1 个数，所以 V[11] = F[11,11]。

1000 = 1010 - 0010，0010 就是十进制的 2，所以 V[1010] 存 2 个数，所以 V[10] = F[10,9]。

0000 = 1000 - 1000，1000 就是十进制的 8，所以 V[1000] 存 8 个数，所以 V[8] = F[8...1]。

V[1011] 存 1 个数， V[1010] 存 2 个数，看的是二进制最右边的一个 1 到末尾的大小。1010 就是 10，1000 就是 1000 。

怎么得到一个数最右边的 1 到末尾的大小，是二进制操作的一个技巧，会用到一些补码的知识，可以参考 趣谈计算机补码。

将原数取反，然后再加 1 得到的新数和原数按位相与就得到了最右边的 1 到末尾的数。

举个例子，对于 101000 ，先取反得到 010111，再加 1 变成 011000，再和原数相与，101000 & 011000，刚好就得到了 1000。其中，取反再加一，根据补码的知识，可以通过取相反数得到。

所以对于 i 的话，i & -i 就得到了最右边的 1 到末尾的数，也就是 V[i] 这个区间存多少个数。

如果 len = i & -i ，那么 V[i] = F[i,i-1,i-2, ... i-len+1]。

参考下边的代码，BIT 就是我们上边要求的 V 数组。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
/**
 * @param {number[]} nums
 */
var NumArray = function (nums) {
    this.nums = [0, ...nums]; //补一个 0
    this.BIT = new Array(this.nums.length);
    for (let i = 1; i < this.BIT.length; i++) {
        let index = i - ( i & -i ) + 1;
        this.BIT[i] = 0;
        //累加 index 到 i 的和
        while (true) {
            this.BIT[i] += this.nums[index];
            index += 1;
            if (index > i) {
                break;
            }
        }
    }
};

有了 BIT 这个数组，一切就都好说了。如果我们想求 F[1, 11] 也就是前 11 个数的和。

F[1,11] = BIT[11] + BIT[10] + BIT[8]，看下二进制 BIT[0001,1011] = BIT[1011] + BIT[1010] + BIT[1000] 。

1011 -> 1010 -> 1000，对于 BIT 每次的下标就是依次把当前数最右边的 1 变成 0 。

这里有两种做法，一种是我们求出当前数最右边的 1 到末尾的数，然后用原数减一下。

举个例子， 1010 最右边的 1 到末尾的数是 10 ，然后用 1010 - 10 就得到 1000 了。

另外一种做法，就是 n & (n - 1)，比如 1010 & (1010 - 1)，刚好就是 1000 了。

知道了这个，我们可以实现一个函数，用来求区间 [1, n] 的和。

1
2
3
4
5
6
7
8
9
NumArray.prototype.range = function (index) {
  let sum = 0;
  while (index > 0) {
    sum += this.BIT[index];
    index -= index & -index;
    //index = index & (index - 1); //这样也可以
  }
  return sum;
};

有了 range 函数，题目中的 sumRange 也就很好实现了。

1
2
3
4
NumArray.prototype.sumRange = function (i, j) {
    //range 求的区间范围下标是从 1 开始的,所以这里的 j 需要加 1
    return this.range(j + 1) - this.range(i);
};

接下来是更新函数怎么写。

更新函数的话，最关键的就是找出，当我们更新数组第 i 个值，会影响到我们的哪些子区间，也就是代码中的 BIT 数组需要更新哪些。

我们来回忆下之前做了什么事情。

这是论文中的一张图，含义就是我们之前分析的，BIT[8] 存的是 F[1...8] ，对应图中的就是从第 8 个位置到第 1 个位置的矩形。BIT[6] 存的是 F[6,5]， 对应图中的就是从第 6 个位置一直到第 5 个位置的矩形。

然后我们水平从某个数画一条线，比如从 3 那里画一条线。

穿过了 3 对应的矩形，4 对应的矩形，8 对应的矩形。因此如果改变第 3 个数，BIT[3]，BIT[4] 以及 BIT[8] 就需要更新。通过这种方式我们把每个数会影响到哪个区间画出来，找一下规律。

当改变了第 5 个元素的时候，会依次影响到 BIT[5]，BIT[6]，BIT[8]，BIT[16]。

00101 -> 00110 -> 01000 -> 10000。

00101 + 1 = 00110。

00110 + 10 = 01000

01000 + 1000 = 10000

可以看到每次都是加上当前数最右边的 1 到末尾的数，即 next = current + (current & -current)。

所以更新的代码也就出来了。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
/**
 * @param {number} i
 * @param {number} val
 * @return {void}
 */
NumArray.prototype.update = function (i, val) {
  i += 1;//对应的下标要进行加 1
  const sub = val - this.nums[i];
  this.nums[i] = val;
  while (i < this.nums.length) {
    this.BIT[i] += sub;
    i += i & -i;
  }
};

综上，这道题就解决了，我们把代码合在一起。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
/**
 * @param {number[]} nums
 */
var NumArray = function (nums) {
  this.nums = [0, ...nums];
  this.BIT = new Array(this.nums.length);
  for (let i = 1; i < this.BIT.length; i++) {
    let index = i - ( i & -i ) + 1;
    this.BIT[i] = 0;
    while (true) {
      this.BIT[i] += this.nums[index];
      index += 1;
      if (index > i) {
        break;
      }
    }
  }
};

/**
 * @param {number} i
 * @param {number} val
 * @return {void}
 */
NumArray.prototype.update = function (i, val) {
  i += 1;
  const sub = val - this.nums[i];
  this.nums[i] = val;
  while (i < this.nums.length) {
    this.BIT[i] += sub;
    i += i & -i;
  }
};
/**
 * @param {number} i
 * @param {number} j
 * @return {number}
 */
NumArray.prototype.sumRange = function (i, j) {
  return this.range(j + 1) - this.range(i);
};

NumArray.prototype.range = function (index) {
  let sum = 0;
  while (index > 0) {
    sum += this.BIT[index];
    // index -= index & -index;
    index = index & (index - 1); //这样也可以
  }
  return sum;
};

/**
 * Your NumArray object will be instantiated and called as such:
 * var obj = new NumArray(nums)
 * obj.update(i,val)
 * var param_2 = obj.sumRange(i,j)
 */

时间复杂度的话，初始化、更新、查询其实都和二进制的位数有关，以查询为例。每次将二进制的最后一位变成 0，最坏的情况就是初始值是全 1，即 1111 这种，执行次数就是 4 次，也就是二进制的位数。

如果是 n ，那么位数大约就是 log(n)，可以结合 再谈进制转换 理解。把一个数展开为 2 的幂次和，位数其实就是最高位的幂次加 1。比如 $11 = 2^0 + 2^1 + 2^3$ ，最高幂次是 3 ，所以 11 的二进制(1011) 位数就是 4。如果要求的数是 n，最高的次幂是 x ，$2^x + … = n$，近似一下 $2^x=n$，x = log(n)，位数就是 log(n) + 1。

所以 update 和 sumRange 的时间复杂度就是 O(log(n))。

对于初始化函数，因为要执行 n 次，所以就是 O(nlog(n))。当然我们也可以利用解法二，把前缀和都求出来，然后更新数组 BIT 的每个值，这样就是 O(n) 了。但不是很有必要，因为如果查询和更新的次数很多，远大于 n 次，那么初始化这里的时间复杂度也就无关紧要了。

总

讲了很多，其实树状数组最根本的就是开头所提到的二进制幂次的分解，$11 = 2^0 + 2^1 + 2^3 = 1 + 2 + 8$，然后把右边的分解出来的数看做子区间的长度。

windliang wechat