MENU

树状数组

April 5, 2019 • 算法

一:前言

树状数组,也叫二叉索引树(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]=\sum_{j=i-LowBit(i)+1}^i a[j] \tag{i ≥ 1}
$$
简单点来说,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[],即,
$$
\begin{align}
b[1] & = a[1] - a[0] \\
b[2] & = a[2] - a[1] \\
... \\
b[n] & = a[n] - a[n-1]
\end{align}
$$
我们来看看引入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]=\sum_{j=i-LowBit(i)+1}^i b[j] \tag{i ≥ 1}
$$
即数组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;
}

五:参考链接

最后编辑于: April 7, 2019 19:39
Archives QR Code Tip
QR Code for this page
Tipping QR Code