树状数组

fortyqin · 2023-08-06 18:34:24 · 算法·理论

引入

我们从这样一个问题引入：给定一个序列 A，有若干次修改和询问。说地详细一些，就是修改序列中的一个数，和查询序列的前缀和。

对于这个问题，我们可以想到使用之前学习过的前缀和，这样查询的时间复杂度为 O(1)，但一次修改操作就要 O(n) 的时间复杂度。

试着折中一下：查询和修改的时间复杂度能否都成为 O(\log_2n)。

反思一下：为什么修改的复杂度这么慢，其实原因就是需要一个一个地修改。那样我们可以想到：我们分成 \log_2n 个组。

我们每次修改还查询只用对这么多个组操作，但是又有一个问题，我们按照什么方式来分组？

与前面学习高效算法相关的，总是和 2 离不开关系。

从一个例子入手

比如我们有一个长度为 8 的数组 a[\;]，要对其进行多次修改和询问，为了分组，我们先写出 1 到 8 每个数的二进制。

我们可以想到，一个十进制数，可以天然地转换为一个 \log_2n 位的二进制数。

再进一步，我们可以想到任意一个是十进制数 n 都可以拆成 m 个不同的 2^i 的数相加，形式化地：

上面的式子，满足 2^{i_1}\ge2^{i_2}\ge...\ge2^{i_{m-1}}\ge2^{i_m}。

对于每一个 2^{i_j}，我们可以让它管理一个区间，就拿 5 来举例：

我们来看怎么计算一个数 x 的 \rm{lowbit}：

首先这个数 x>0，设 x 的第 k 位是 1，并且所有低于 k 的位数都为 0。

将 x 按位取反，此时第 k 位是 0，所有低于 k 的位数都为 1。

此时再给 x+1，此时第 k 位是 1，所有低于 k 的位数都为 0，而高于 k 的数都和原来的数值不同。

此时再按位与上最开始的 x，此时只有第 k 位是 1。

因为 x 的取反 +1 就是负数的表示法，所以可以按照上面的式子简写。实现如下：

int lowbit(int x) { // 查询 x 的 lowbit
    return x & (-x);
}

初始化

对于原始数组 a[\;]，我们使 c[\;] 初始都为 0（全局数组），我们该怎么赋值呢？

其实非常简单，我们可以直接调用后面的单点修改函数，这样就可以解决这个问题，时间复杂度 O(n\log_2n)。

当然还有更高效的方法，我们暂时先不讨论。

单点修改+区间查询

例题：【模板】树状数组 1

首先树状数组要实现的第一个功能是查询前缀和。根据前面的讨论，我们知道 [1,x] 区间在 c[\;] 中，最多分成 \log_2x 个部分。

我们可以每次访问 c_i 并累计到答案中，然后把 i 加上 \rm{lowbit}(i)，实现如下：

int query(int x) { // 查询区间 [1,x] 的和
    int ret = 0;
    for (int i=x; i<=n; i+=lowbit(i)) { // n 是区间长度
        ret += c[i];
    }
    return ret;
}

当然知道前缀和我们就可求得区间和，那么就不在赘述。

树状数组的第二个功能是单点修改，其实也非常简单，将包括 x 的 c[i] 都统一加上一个 k。

我们可以每次访问 c_i 并将其加上 k，然后把 i 减去 \rm{lowbit}(i)，实现如下：

void updata(int x, int k) { // 给 a[x] 加上 k
    for (int i=x; i<; i-=lowbit(i)) {
        c[i] += k;
    }
}

小小的空间优化

因为其实代码我们根本就没有使用 a[\;]，其实可以直接在 c[\;] 上进行操作。

区间修改+单点查询

例题：【模板】树状数组 2

我们知道差分是前缀和的逆运算，我们可以使 c[\;] 为一个差分数组。

这样可以很方便地实现这个功能，因为代码和上面的相同，只是调用时有所不同，所以不再给出代码。

但还有一个小细节：初始时的 a[\;] 经过差分后才能初始化 c[\;]。

区间修改+区间查询

例题：【模板】线段树 1

虽然这是线段树的模板，树状数组依然可以做。

我们依然让 c[\;] 为差分数组，我们重点讨论区间查询。求 [1,n] 的和，那么式子为：

看着很繁琐，不过没关系，我们将它展开，看看每个 c_j 被加了几次，那么得到式子为：

接着将它拆分开来：

此时，我们只需要维护两个树状数组：c[\;] 和 d[\;]，其中 d_i=c_i\times i。

我们给出代码细节，注意添加注释的地方：

int query(int x) { // 查询区间 [1,x] 的和
    int ret = 0;
    for (int i=x; i<=n; i+=lowbit(i)) { 
        ret += (x+1)*c[i]-d[i]; // 注意是 (x+1) 不是 (i+1)
    }
    return ret;
}
void updata(int x, int k) { // 给 a[x] 加上 k    
    for (int i=x; i<; i-=lowbit(i)) {        
        c[i] += k; 
        d[i] += x*k; // 注意 x
    }
}

值域树状数组

逆序对

例题：逆序对

读完题目后，我们倒着来想。

对于位置 i 我们看 [i+1,n] 有多少比它小的，也就是说我们找后面值域在 [1,a_i-1] 的数有几个。

这样，我们树状数组可以维护：值域在 [1,x] 之间数的个数，这借用了计数排序的思想。

当然要注意：本题需要离散化，基本上很多值域树状数组的题目都需要离散化。

这里给的代码是调用时的片段：

for (int i=n; i>=1; i--) {
    updata(a[i], 1); // 给值为 a[i] 的地方数量 +1
    ans += sum(a[i]-1); // 统计此时逆序对的个数
}

普通平衡树

例题：【模板】普通平衡树

其实这是一个值域树状数组也可以做的题目，我们依次来实现需要的功能。

• 插入 x 数，删除 x 数。

  这是基础操作，不再赘述。

• 查询 x 数的排名（排名定义为比当前数小的数的个数 +1）。

  那么直接求小于 x 数的个数再 +1，就可以实现。

• 查询排名为 X 的数。

  由于不会出现负数，那么 c[\;] 的前缀和是单调递增的，我们可以进行二分查找。

• 求 x 的前（后）驱（前驱定义为小于 x，且最大的数；后继定义为大于 x，且最小的数）。

  看着很难，其实很简单，只需要先求出 x 的排名，再求出排名为 x\pm1 的数。

这道题目还有一个细节，值域会有负数，那么整体偏移成正数就可以。

二维：单点修改+区间查询（区间修改+单点查询）

例题：二维树状数组 1：单点修改，区间查询

例题：二维树状数组 2：区间修改，单点查询

前置知识：二维差分+二维前缀和，默认你已经掌握了。

那么其实很简单，只需要再多出来一维，成为 c[\;][\;] 就可以了，下面是代码：

int query(int x, int y) {
    int ret = 0;
    for (int i=x; i; i-=lowbit(i)) {
        for (int j=y; j; j-=lowbit(j)) {
            ret += c[i][j]; // 照常求和
        }
    }
    return ret;
}
void updata(int x, int y, int k) {
    for (int i=x; i<=n; i+=lowbit(i)) {
        for (int j=y; j<=m; j+=lowbit(j)) { // 注意 m
            c[i][j] += k; // 照常修改
        }
    }
}

学会这些后，你区间修改+单点查询也学会了。

多维：区间修改+区间查询

前置知识：二（三）维差分+二（三）维前缀和。

和一维的区间修改+区间查询一样，都是维护多个树状数组，并且树状数组本身是个差分数组。

那么接下来只是推式子，并不给出具体的实现代码。

当然区间修改和普通的一样，重点推如何计算前缀和。

二维式子

我们来考虑一个 c_{i,j} 在一个左上角为 (1,1) 右下角为 (x,y) 的矩阵中出现的次数。

可以看到只有在一个左上角为 (i,j) 右下角为 (x,y) 矩阵中的每个数计数前缀和是要调用 c_{i,j} 的。

所以，只需要求这个矩阵中数的个数，就知道每个 c_{i,j} 被调用了几次。

我们列出式子并进行拆分：

最后我们得出需要 4 个树状数组可以搞定二维区间修改+区间查询。

三维式子

同二维一样，我们要考虑每个 c_{i,j,k} 出现的次数。

此时我们再增加一维高度，只有高度大于等于它的点才会调用它。

形式化地，只有满足 a\ge i 且 b \ge j 且 c\ge k 时，c_{a,b,c} 才会调用 c_{i,j,k}（下标 c，不是数组 c[\;][\;][\;] ）。

我们列出式子并进行拆分，可以自行手算，这里直接给出最终的结果：

最后我们得出需要 8 个树状数组可以搞定三维区间修改+区间查询。

树上树状数组

对于树上问题，我们不能直接做，但是我们可以将其转换为序列上的问题，这样就很容易。

DFS 序

简单来说，我们对一棵树进行 DFS，记录每个点第一次访问的时间，和每个子树的大小。

这样一颗以 x 为根的子树对应的区间是：[\rm{dfn}(x),\rm{dfn}(x)+size(x)-1]。

其中 \rm{dfn}(x) 代表 x 第一次访问的时间，\rm{size}(x) 代表以 x 为根子树的大小。

这样我们可以解决以下的问题：

• 单点修改+子树查询
• 子树修改+单点查询
• 子树修改+子树查询

这和之前解决的问题类似，也就只给出遍历树的代码：

int tot = 0;
void dfs(int x, int fa) { // 遍历到节点 x，它的父亲是 fa
    dfn[x] = ++tot; // 使用全局变量 tot 记录
    size[x] = 1;
    for (int i=head[x]; i; i=edge[i].next) { // 链式前向星
        int y = edge[i].to; // y 是 x 的儿子
        if (y == fa) continue; // 双向边判重
        dfs(y, x);
        size[x] += size[y];
    }
}

树上差分

前置知识：LCA（最近公共祖先）。

主要讨论点差分，简单来说：一个树中，x 节点的点权，表示为在差分树中 x 为根的子树的点权之和。

那我们来解决链修改+子树查询的问题。

其实非常简单我们先给链的两个端点 x 和 y 都加上修改的值 k。

此时 \rm{lca}(x, y) 修改了两次，我们给它减去一个 k。

但是还没完，\rm{father}(\rm{lca}(x, y)) 是不用修改的，我们再给它减去一个 k，其中 \rm{father}(x) 代表 x 的父亲。

总结

其实，写了这么多，有没有发现其实树状数组最终都将问题转为求前缀和，由此求得区间。

其实这里和可以理解为是指一种满足结合律的运算的累算，而如果这种计算不太好表示为前缀和的形式，那么树状数组会非常不方便。