Duanyll's Blog

树状数组的奇妙运用

07 October 2019 | duanyll | Tags: OI 数据结构

树状数组是一种能够$O(\log n)$在线维护前缀和的数据结构, 其写法简单常数小… 不具体介绍了, 看看一些奇妙的操作吧.

人人都会的

  • 单点修改, 区间查询: 基本维护前缀和
  • 区间修改, 单点查询: 维护差分

区间修改, 区间查询

引入数组$d$, $d_i$表示区间$[i, n]$中需要加值的差分, 进行区间加法时, 就直接对$d$操作, 对$d_l$加上$x$, 对$d_{r+1}$减去$x$. 查询前缀和时, 设$sum_i$为区间$[1,i]$的元素和, 易得

\[\begin{aligned} sum_i &= \sum_{j=1}^i(a_j) + \sum_{j=1}^i\sum_{k=1}^j(d_k) \\ &= \sum_{j=1}^i(a_j) + \sum_{j=1}^i(d_j\times(i-j+1)) \\ &= \sum_{j=1}^i(a_j) + (i+1)\sum_{j=1}^i(d_j) + \sum_{j=1}^i(d_j*j) \end{aligned}\]

发现我们需要维护$a_i$, $d_i$, $d_i\times i$的前缀和. 代码可以这样写:

int suma[MAXN], sumd[MAXN], sumdi[MAXN];

void update(int* c, int pos, int x) {
    while (pos <= n) {
        c[pos] += x;
        pos += lowbit(pos);
    }
}

void sum(int* c, int pos) {
    int res = 0;
    while (pos >= 1) {
        res += c[pos];
        pos -= lowbit(pos);
    }
}

void add(int l, int r, int x) {
    update(sumd, l, x);
    update(sumd, r + 1, -x);
    update(sumdi, l, l * x);
    update(sumdi, r + 1, -(r + 1) * x);
}

int query(int l, int r) {
    return suma[r] + r * sum(sumd, r) + sum(sumdi, r)
           - (suma[l - 1] + (l - 1) * sum(sumd, l - 1) + sum(sumdi, l));
}

二维前缀和

首先, 由二维前缀和求矩阵$(x_1, y_1)\rightarrow(x_2,y_2)$的操作方法:

\[sum=sum[x_2][y_2]-sum[x_1-1][y_2]-sum[x_2][y_1-1]+sum[x_1-1][x_2-1]\]

用树状数组维护二维前缀和的方法:

void update(int x, int y, int val) {
    while (x <= n) {
        int t = y;
        while (t <= m) {
            c[x][t] += val;
            t += lowbit(t);
        }
        x += lowbit(x);
    }
}

int query(int x, int y) {
    int res = 0;
    while (x >= 1) {
        int t = y;
        while (t >= 1) {
            res += c[x][t]
            t -= lowbit(t);
        }
        x -= lowbit(x);
    }
    return res;
}

因此不难做到二维单点修改, 矩阵查询.

二维矩阵修改, 矩阵查询

首先说明一下二维差分怎么写. 由二维前缀和的形式, 定义

\[d[i][j] = a[i][j] - a[i-1][j] - a[i][j-1] + a[i-1][j-1]\]

因此

\[a[i][j] = \sum_{n=1}^i\sum_{m=1}^j(d[n][m])\]

即可用二维树状数组维护差分.

比如给一个$5\times5$的矩阵中间$3\times3$的部分加上$x$, 可以这么修改差分数组:

 0  0  0  0  0
 0 +x  0  0 -x
 0  0  0  0  0
 0  0  0  0  0
 0 -x  0  0 +x

所以矩阵修改操作的写法:

void add(int x1, int y1, int x2, int y2, int x) {
    update(x1, y1, x);
    update(x1, y2 + 1, -x);
    update(x2 + 1, y1, -x);
    update(x2 + 1, y2 + 1, x);
}

为了实现二维区间修改和查询的操作, 我们可以类比一维的情况, 用差分来表示前缀和.

\[\begin{aligned} sum[x][y] &= \sum_{i=1}^x\sum_{j=1}^y(a[i][j]) \\ &= \sum_{i=1}^x\sum_{j=1}^y\sum_{n=1}^i\sum_{m=1}^j(d[n][m]) \end{aligned}\]

根据前面的套路, 我们统计每一个$d[i][j]$的出现次数, 可以展开两层$\Sigma$: \(\begin{aligned} sum[x][y] =& \sum_{i=1}^x\sum_{j=1}^yd[i][j] * (x-i+1) * (y-j+1) \\ =& (x+1)*(y+1)\sum_{i=1}^x\sum_{j=1}^yd[i][j] \\ & -(y+1)\sum_{i=1}^x\sum_{j=1}^yd[i][j]*i-(x+1)\sum_{i=1}^x\sum_{j=1}^yd[i][j]*j \\ & +\sum_{i=1}^x\sum_{j=1}^yd[i][j]*i*j \end{aligned}\)

也就是说, 我们需要维护

  • $\sum_{i=1}^x\sum_{j=1}^yd[i][j]$
  • $\sum_{i=1}^x\sum_{j=1}^yd[i][j]*i$
  • $\sum_{i=1}^x\sum_{j=1}^yd[i][j]*j$
  • $\sum_{i=1}^x\sum_{j=1}^yd[i][j]ij$

开四个树状数组即可. 这样用树状数组写, 常数会比线段树小一些.