树状数组 - 理解至上

oldSix
843 阅读5分钟
原文链接： subetter.com

一：前言

树状数组，也叫二叉索引树（Binary Indexed Tree，简称 BIT），由 Peter M. Fenwick 于 1994 年发表。其初衷是解决数据压缩里累积频率（Cumulative Frequency）的计算问题，现多用于动态地计算数组的区间和。

二：单点更新，区间查询

先看一个简单的问题：有一个长度为 1000，初始为 0 的整型数组，不妨设为int a[1000]，现对它有两种操作，

  1. 给某位置的元素加上delta；
  2. 求某个区间内的所有元素的总和。

树状数组的做法是，

#include <iostream>

using namespace std;

#define N 1000

int a[N];
int tree[N];

/* 除了最低位，清空其它所有为 1 的位 */
int LowBit(int x)
{
    return x & (-x);
}

/* 元素 a[i] 的加 delta */
void Add(int i, int delta)
{
    // 因为 tree[i] 正好可以覆盖 a[i]，所以此句代码包括数组 a[] 的声明定义
    // 皆可省略，加上它们只为读者方便理解
    a[i] += delta;

    while (i <= N)
    {
        tree[i] += delta;
        i += LowBit(i);
    }
}

/* 求数组 a[] 的前 i 项和 */
int Sum(int i)
{
    int sum = 0;

    while (i > 0)
    {
        sum += tree[i];
        i -= LowBit(i);
    }

    return sum;
}

/* 求区间 a[i, j] 的和 */
int RangeSum(int i, int j)
{
    return Sum(j) - Sum(i - 1);
}

int main()
{
    memset(a, 0, sizeof(a));
    memset(tree, 0, sizeof(tree));

    Add(3, 2);
    cout << RangeSum(1, 3) << endl; // 2
    cout << RangeSum(2, 4) << endl; // 2

    Add(10, 1);
    cout << RangeSum(1, 3) << endl; // 2
    cout << RangeSum(5, 12) << endl; // 1
    cout << RangeSum(1, 11) << endl; // 3

    return 0;
}

三：算法分析

我们先给出上述代码中tree[i]的定义，
tree[i]=i∑j=i−LowBit(i)+1a[j]
简单点来说，tree[i]表示数组a[]中连续LowBit(i)个元素的和。我们配合上图，举几个例子，

  • 当 i = 2 时，LowBit(2) = 2，则tree[2] = a[2] + a[1]
  • 当 i = 5 时，LowBit(5) = 1，则tree[5] = a[5]
  • 当 i = 8 时，LowBit(8) = 8，则tree[8] = a[8] +a[7] + ... + a[1]
  • 当 i = 12 时，LowBit(12) = 4，则tree[12] = a[12] + a[11] + a[10] +a[9]

从图中很容易看出，每一个tree[i]结点，都用一个顶部带有灰色的矩形表示，而它的高度正好等于它所表示的范围。由tree[i]的定义，我们可以很容易得出两条性质，

性质一：Add()中从任意位置 i（i > 0）往上进行更新（利用i += LowBit(i)），范围可以覆盖a[i]的所有tree[]结点都能得到更新。

性质二：Sum()中从任意位置 i （i > 0）往下求和（利用i -= LowBit(i)），经过的所有的tree[]结点，它们所表示的区间范围，都能恰好地衔接在一起（即区间既不会重叠，也不会断开）。

现在我们假设这样的场景：先对a[i]进行更新，即进行Add(i, delta)操作，那么在向上更新所经过的tree[]结点，我们把它们标记为绿色，并用线依次连接；接下来求和Sum(j)，其中 0 < i < j，我们再把这时候向下经过的tree[]结点标记为橙色，也用线依次连接。

想一想，这两条颜色不同的链表会是什么样子？

这就是性质三的内容，即下图所示，

性质三：Add(i, delta)和Sum(j)所经过的tree[]结点，有且仅有一个结点是相同的，其中 0 < i < j。

证明如下：

根据 0 < i < j，从二进制角度，高位开始，找到第一个不同的位，

由于i += LowBit(i)和j -= LowBit(j)，所以它们一定会存在相等的情况。

对于唯一性的证明，用反证法即可。

假设存在两个或两个以上相同的结点，取其中最小的结点，我们发现，分别按照i += LowBit(i)和j -= LowBit(j)的更新规律，是无法得到其它较大的结点的，因此假设失败，只会存在一个相同的结点。

四：区间更新，单点查询

再来看一个经典的问题。

有一个长度为 1000，初始为 0 的整型数组，不妨设为int a[1000]，现对它有两种操作，

  1. 区间 [i, j]内的元素分别加上delta；
  2. 求a[i]的值。

如何用树状数组来解决这个问题呢？

引入a[]的差分数组，设其为b[]，即，
b[1]=a[1]−a[0]b[2]=a[2]−a[1]...b[n]=a[n]−a[n−1]
我们来看看引入b[]的好处。

对于第一个操作：区间 [i, j]内的元素分别加上delta，相当于b[i] += delta; b[j + 1] -= delta。

对于第二个操作：求a[i]的值，则相当于求和 a[i] = b[i] + b[i-1] +...+b[1]。

由此，我们定义tree[i]，
tree[i]=i∑j=i−LowBit(i)+1b[j]
即数组b[]中连续LowBit(i)个元素的和。

代码实现如下，

#include <iostream>

using namespace std;

#define N 1000

int tree[N];

/* 除了最低位，清空其它所有为 1 的位 */
int LowBit(int x)
{
    return x & (-x);
}

/* a[i, N] 的值各加 delta，相当于 b[i] 加 delta */
void Add(int i, int delta)
{
    while (i <= N)
    {
        tree[i] += delta;
        i += LowBit(i);
    }
}

/* a[i, j] 的值各加 delta，相当于 b[i] 加 delta，b[j + 1] 减 delta */
void RangeAdd(int i, int j, int delta)
{
    Add(i, delta);
    Add(j + 1, -delta);
}

/* 求 a[i] */
int Query(int i)
{
    int sum = 0;

    while (i > 0)
    {
        sum += tree[i];
        i -= LowBit(i);
    }

    return sum;
}

int main()
{
    memset(tree, 0, sizeof(tree));

    RangeAdd(2, 5, 1);
    cout << Query(1) << endl; // 0
    cout << Query(3) << endl; // 1
    cout << Query(6) << endl; // 0

    RangeAdd(3, 8, 2);
    cout << Query(2) << endl; // 1
    cout << Query(4) << endl; // 3
    cout << Query(10) << endl; // 0

    return 0;
}

五：参考文献

avatar
文章
阅读
粉丝