树上启发式合并学习笔记

最近几天了解到一个很神奇的算法——dsu on tree，看上去没多快实际上很快，这叫低调。

好久不更了，至于反演，5 月再更吧，4 月的最后一天分享一下 dsu on tree。顺便闲话一句，4/26 是我生日，也是历史二模。

重链剖分 dsu on tree

这类 dsu on tree 适用于多次询问，每次询问需要 $O(\mathrm{子}\mathrm{树}\mathrm{大}\mathrm{小})$ 的时间，这时用重链剖分 dsu on tree 可以优化到 $O(q\log n)$

例题1：

CF600E，典中典了。

给一棵以 $1$ 为根的树，每个点都有一个颜色，求出以每个点为根的子树中主导颜色的编号之和。

所谓主导颜色，就是出现次数最多的颜色，可能有多个。

考虑一个朴素的做法：

开一个桶存储每个颜色的出现次数，按照 DFS 序遍历，遍历到一个点时，依次统计子结点的答案，统计完一个就把桶清空。

接着再把它子结点的颜色加到桶里，顺便处理出以当前节点为子树的主导颜色。

这样的做法为 $O(n\times n)$，超时。

 

但是可以优化，我们定义一个非叶子节点的重儿子是 $size$ 最大的儿子，轻儿子是除了重儿子外的儿子。

对于一个节点 $x$，最后一个儿子算完之后是不用清空桶的，它可以直接记录到以 $x$ 为根的子树的答案中，剩下的依旧是暴力。

显然把最大的一个儿子设为重儿子最优了，优化成了 $O(n\log n)$。

 

下面来证明一下，每个节点被统计的次数为它上面的轻边数量（如果是重边，统计完就会给它的父亲了）。

每经过一个轻边，它的兄弟中就一定有一个 size 比它大的。

所以经过后子树的 $size$ 至少翻倍，$\log n$ 次 到达根，每个节点被统计 $\log n$ 次，$n$ 个就被统计 $n\log n$ 次了，证毕。

 

注意：清空桶时，如果全清空还是超时，所以只需要清空用过的就行了。

这里使用 $map$，虽然时间复杂度会多一个 $\log$，但是一句 m.clear() 就行了还很快。

 

代码（$n{\log }^{2}n$ 的）：

#include <map>
#include <vector>
#include <iostream>
#define int long long
using namespace std;
int n, ma, res;
int fa[100005], sz[100005], wson[100005];
int c[100005], ans[100005];
vector <int> v[100005];
map <int, int> m, b;
int dfs1 (int x) {
    for (int i = 0; i < v[x].size (); i ++) if (v[x][i] != fa[x]) {
        fa[v[x][i] ] = x;
        int tmp = dfs1 (v[x][i]);
        sz[x] += tmp;
        if (sz[wson[x] ] < tmp) wson[x] = v[x][i];
    }
    return ++ sz[x];
}
void add (int x) {
    m[x] ++;
    if (m[x] > ma) {
        ma = m[x];
        res = x;
    } else if (m[x] == ma) res += x;
}
void dfs2 (int x) {
    add (c[x]);
    for (int i = 0; i < v[x].size (); i ++) if (v[x][i] != fa[x]) dfs2 (v[x][i]);
}
void dfs_ans (int x) {
    for (int i = 0; i < v[x].size (); i ++) {
        if (v[x][i] != fa[x] && v[x][i] != wson[x]) {
            dfs_ans (v[x][i]);
            m.clear ();
            res = ma = 0;
        }
    } 
    if (wson[x] != 0) dfs_ans (wson[x]);
    add (c[x]);
    for (int i = 0; i < v[x].size (); i ++) if (v[x][i] != fa[x] && v[x][i] != wson[x]) dfs2 (v[x][i]);
    ans[x] = res;
}
signed main () {
    int x, y;
    scanf ("%lld", &n);
    for (int i = 1; i <= n; i ++) cin >> c[i];
    for (int i = 1; i < n; i ++) {
        scanf ("%lld%lld", &x, &y);
        v[x].push_back (y);
        v[y].push_back (x);
    }
    dfs1 (1);
    dfs_ans (1);
    for (int i = 1; i <= n; i ++) printf ("%lld ", ans[i]);
    return 0;
}

 例题2：

CF375D

关于询问，我们用一个 vector<pair<int, int> > v[100005] 来记录，v[i] 存的是询问中的第一个参数为 i 的。

pair 的第一个是 $k$，第二个是询问编号，每遍历到一个点，处理完信息后，访问 v[x] 并把询问搞掉。

对于这题，可以用一个桶和一个 $BIT$ 来维护。

桶维护的是颜色出现的次数，$BIT$ 维护的是颜色出现次数的出现次数。

先来考虑询问，假设当前的桶和 $BIT$ 已经维护好了，要求出现次数 $\ge k$ 的颜色数量，query (100000) - query (k - 1) 即可。

每次把 $BIT$ 和桶清空，还是最后一个遍历重儿子。

这里为了方便，两者都用 $map$ 写，所以时间复杂度是 dsu on tree 的一个 $\log$，$BIT$ 的一个 $log$，$map$ 的一个 $\log$。

$n{\log }^{3}n$，但是凭借 $BIT$ 和 dsu on tree的超小常数，过了。

代码：

#include <map>
#include <vector>
#include <iostream>
using namespace std;
int n, q;
int fa[100005], sz[100005], wson[100005];
int c[100005], ans[100005];
vector <pair <int, int> > v[100005];
vector <int> G[100005];
map <int, int> sum, b;//sum: 树状数组，维护每个数出现次数的出现次数
int dfs1 (int x) {//剖链
    for (int i = 0; i < G[x].size (); i ++) if (G[x][i] != fa[x]) {
        fa[G[x][i] ] = x;
        int tmp = dfs1 (G[x][i]);
        sz[x] += tmp;
        if (sz[wson[x] ] < tmp) wson[x] = G[x][i];
    }
    return ++ sz[x];
}
void add (int x, int y) {for (; x <= 100000; x += x & -x) sum[x] += y;}
int query (int x) {
    int ret = 0;
    for (; x > 0; x -= x & -x) ret += sum[x];
    return ret;
}
void dfs2 (int x) {//重新处理，之前的被清空了。 
    if (b[c[x] ] != 0) add (b[c[x] ], -1);
    ++ b[c[x] ];
    add (b[c[x] ], 1);
    for (int i = 0; i < G[x].size (); i ++) if (G[x][i] != fa[x]) dfs2 (G[x][i]);
}
void dfs_ans (int x) {//统计以 x 为根的子树的答案。 
    for (int i = 0; i < G[x].size (); i ++) if (G[x][i] != fa[x] && G[x][i] != wson[x]) {
        dfs_ans (G[x][i]);
        sum.clear ();
        b.clear ();
    }
    if (wson[x] != 0) dfs_ans (wson[x]);
    if (b[c[x] ] != 0) add (b[c[x] ], -1);
    ++ b[c[x] ];
    add (b[c[x] ], 1);
    for (int i = 0; i < G[x].size (); i ++) if (G[x][i] != fa[x] && G[x][i] != wson[x]) dfs2 (G[x][i]);
    for (int i = 0; i < v[x].size (); i ++)
        ans[v[x][i].second] = query (100000) - query (v[x][i].first - 1);
    //second 是询问编号，first 是 k。出现次数在 1 ~ 100000 的减去在 1 ~ k - 1 的得到在 1 ~ k 的。 
}
int main () {
    int x, y;
    scanf ("%d%d", &n, &q);
    for (int i = 1; i <= n; i ++) scanf ("%d", &c[i]);
    for (int i = 1; i < n; i ++) {
        scanf ("%d%d", &x, &y);
        G[x].push_back (y);
        G[y].push_back (x);
    }
    dfs1 (1);
    for (int i = 1; i <= q; i ++) {
        scanf ("%d%d", &x, &y);
        v[x].push_back (make_pair (y, i) );
    }
    dfs_ans (1);
    for (int i = 1; i <= q; i ++) printf ("%d\n", ans[i]);
    return 0;
}

长链剖分 dsu on tree

重链剖分 dsu on tree 的题目还有很多，这里不再详解了。长链剖分才是优美中的优美，以 $O(n)$ 解决问题。

它适用于解决这类问题：多次询问，每次询问也是求以 $x$ 为根的子树内的东西，而直接暴力的复杂度是 $depth[x]$ 的。（即以 $x$ 为根的子树的深度）

例题：

CF1009F

题意十分简单，对于每个结点，求它子树内同一深度结点之和的最大值，显然可以线段树合并，这个我改天再讲。

大家已经看出这个东西可以用重链剖分 dsu on tree 解决了吧，但是加上 $map$ 的一个 log 就有点难过了，所以考虑长链剖分 dsu on tree。

还是先考虑朴素做法，先给每个点赋值一个深度值，这个值就是以 $1$ 号点为根的深度，到时候算距离了减下当前的深度即可。

然后设 $f[i][j]$ 为以 $i$ 为根的子树中，深度为 $j$ 的结点个数，转移很简单，如下：

int ans = 0;
for (int son : G[i]) {
    for (int j = dep[son]; j <= max_dep[son]; j ++) {//< dep[son] 的没有节点，> max_dep[son] 的也没有节点。
        //注意这里的 dep 指的是以 1 号节点为根的深度。
        f[i][j] += f[son][j];
        if (f[i][j] >= f[i][ans]) ans = j;
    }
}
f[i][dep[i] ] ++;

这样把一个儿子合并到父节点的时间复杂度就是以这个儿子为根的子树的深度。

时间复杂度为 $O(\sum max$_$dep[i])$，当它是接近满二叉树时，时间复杂度为 $O(n+2\times n/2+4\times n/4+...)=n\log n$ 的。

如果是接近一条链，那么直接爆炸。

所以来考虑如何优化，显然最后一个儿子是不用算的，直接送给父节点即可。

然后我们发现把最长的儿子设为最后一个十分的优秀，以最长的儿子为根的深度是以父节点为根的深度减去 $1$。

所以时间复杂度变为 $O(\sum max$_$dep[i]-\sum (max$_$dep[i]-1))=O(n)$ 了，并且空间复杂度也是 $O(n)$。

神奇吧！指针写的，没用 $map$，省了一个 log：

#include <iostream>
#include <vector>
const int maxn = 1001001;
std::vector<int> G[maxn];
int depth[maxn];
int ans[maxn];
int son[maxn];

void dfs1(int u, int fa) {
    depth[u] = 1;
    for (auto v : G[u]) 
      if (v != fa) {
           dfs1(v, u);
           depth[u] = std::max(depth[u], depth[v] + 1);
           if (son[u] == 0 or depth[son[u]] < depth[v])
              son[u] = v;
      }
}

int *D[maxn];
int buf[maxn];
int *head = buf;

void dfs2(int u, int fa) {    
    D[u][0] = 1;
    ans[u] = 0;
    if (son[u] > 0) {
        // first visit longest son        
        D[son[u]] = D[u] + 1;
        dfs2(son[u], u);
        
        if (D[son[u]][ans[son[u]]] > 1) 
           ans[u] = ans[son[u]] + 1; 
        
        // second visit other son
        for (auto v : G[u]) 
          if (v != fa and v != son[u]) {
              D[v] = head;
              head += depth[v];
              dfs2(v, u);
              for (int j = 0; j < depth[v]; j++) {
                 D[u][j+1] += D[v][j];
                 if (D[u][j+1] > D[u][ans[u]])
                    ans[u] = j + 1;
                 else if (D[u][j+1] == D[u][ans[u]] and j + 1 < ans[u])
                    ans[u] = j + 1;
            }
          }
    }
}

int main() {
    int n;
    scanf ("%d", &n);
    for (int i = 1; i < n; ++i) {
        int x, y;
        scanf ("%d%d", &x, &y);
        G[x].push_back(y);
        G[y].push_back(x);
    }
    dfs1(1, 0);
    D[1] = head;
    head += depth[1];
    dfs2(1, 0);
    for (int i = 1; i <= n; ++i)
      printf ("%d\n", ans[i]);
}

 

结语

关于 dsu on tree 就介绍到这里了，这样的题目虽然很少，但要是考到了不会就尴尬了。