树上启发式合并 学习笔记

· · 算法·理论

树上启发式合并(DSU on tree)

引入 CF375D Tree and Queries

一颗以 1 为根的树,每个顶点都有一个颜色 c_{v}
对于每个询问 v_{j},k_{j},你需要回答在以 v_{j} 为根的子树中,有多少种颜色在该子树中出现的次数 \geq k_{j}

分析

暴力

首先离线下所有询问。
对于每个结点,我们 DFS 其子树,
并用桶标记颜色 x 的出现次数 sum_x,以及出现次数大于等于 i 的颜色数 cnt_i
所以每个子结点 v 的贡献为:

sum[c[v]] ++;
cnt[sum[c[v]]] ++;
//因为 sum[] 是从 0 累加上来的
//所以 cnt[i] 中 i <= sum[c[v]] 的贡献也被计入了,是正确的

遍历完子树后,结点 u 的答案就已经被算好了,然后处理所有关于 u 的询问并记录就行了。

下一步,清空桶和 cnt(因为这个结点的子树信息和下一个结点是不一样的)。
重复上面的步骤把 n 个结点都算一遍就行了。 ::::info[Code]

#include<iostream>
#include<vector>
#include<cstring>

using namespace std;

const int N = 1e5+5;
int n,m,c[N];
vector<int> g[N];

struct Query {
    int id,k;
};
vector<Query> q[N];//询问 
int ans[N];//第 i 次询问的答案 

int sum[N],cnt[N];
void init() {
    memset(sum,0,sizeof(sum));
    memset(cnt,0,sizeof(cnt));
}
void dfs_solve(int u,int fa) {
    sum[c[u]] ++;
    cnt[sum[c[u]]] ++;
    for(auto v:g[u]) {
        if(v == fa) continue;
        dfs_solve(v,u);
    }
}
void solve(int u,int fa) {
    dfs_solve(u,fa);
    for(auto t:q[u]) {
        ans[t.id] = cnt[t.k];
    }
    init();
}

void dfs(int u,int fa) {
    solve(u,fa);
    for(auto v:g[u]) {
        if(v == fa) continue;
        dfs(v,u);
    }
}

int main() {
    cin >> n >> m;
    for(int i = 1; i <= n; i ++)
        cin >> c[i];
    for(int i = 1; i < n; i ++) {
        int u,v;
        cin >> u >> v;
        g[u].push_back(v);
        g[v].push_back(u);
    }
    for(int i = 1; i <= m; i ++) {
        int v,k;
        cin >> v >> k;
        q[v].push_back({i,k});
    }
    dfs(1,0);
    for(int i = 1; i <= m; i ++) {
        cout << ans[i] << '\n';
    }
    return 0;
}

::::

时间复杂度 O(n^2)

树上启发式合并

上面的暴力,每次都要清空,太浪费了。
考虑优化。

暴力计算完结点 u 的答案后,应计算其子结点 v 的答案。
但其实,v 的信息我们已经处理过一遍了(计算 u 答案时),因此可以继承。

问题又来了,每个子结点的子树信息都不一样,
所以算完 v_1 再算 v_2 时还要把得到的信息清空,
因而只有最后处理的一个子结点的信息可以被保留。

很显然,要保留携带的信息最多的子结点(即重子结点),因为可以最好的优化。

::::info[重子结点] 在树链剖分/重链剖分中用过。

表示其子结点中子树最大的子结点。

int sz[N],son[N];
//记录子树大小和重子结点
void dfs(int u,int fa) {
    sz[u] = 1;
    for(auto v:g[u]) {
        if(v == fa) continue;
        dfs(v,u);
        if(sz[v] > sz[son[u]]) son[u] = v;
        sz[u] += sz[v];
    }
}

::::

实现

遍历到结点 u,我们按以下的步骤处理:

时间复杂度

看起来,这样操作的时间复杂度没有优化,
但可以证明降到了 O(n\log n)

细节处理

维护子树方法 — DFS 序

在删改子树信息时,我们可以借助 DFS 序快速遍历子树结点。
常数较小。 ::::info[DFS 序] 深度优先搜索一棵树时,每个结点第一次被遍历到的顺序。

上图的 DFS 序为 1 2 3 4 5 u v 8 9 10

在 DFS 序中,同一颗子树的结点编号是连续的,
所以可以将子树信息转为线性的区间信息,
配合线段树或树状数组进行维护。

$sz_u$ 表示结点 $u$ 的子树大小。 所以 DFS 序的区间 $[dfn_u,dfn_u+sz_u-1]$ 为 $u$ 的子树。 ::::

上文例题代码

#include<iostream>
#include<vector>

using namespace std;

const int N = 1e5+5;
int n,m,c[N];
vector<int> g[N];

int dfn[N],rnk[N],sz[N],son[N],tot = 0;
void dfs1(int u,int fa) {
    dfn[u] = ++tot;
    rnk[tot] = u;
    sz[u] = 1;
    for(auto v:g[u]) {
        if(v == fa) continue;
        dfs1(v,u);
        if(sz[v] > sz[son[u]]) son[u] = v;
        sz[u] += sz[v];
    }
}

struct Query {
    int id,k;
};
vector<Query> q[N];//询问 
int ans[N];//第 i 次询问的答案 

int sum[N],cnt[N];
inline void add(int u) {
    x = c[u];
    sum[x] ++;
    cnt[sum[x]] ++;
}//添加贡献 
inline void del(int u) {
    x = c[u];
    cnt[sum[x]] --;
    sum[x] --;
}//删除贡献 

void dfs2(int u,int fa,bool f) {
//f 表示当前结点是否为轻子结点 
    for(auto v:g[u]) {
        if(v == fa || v == son[u]) continue;
        dfs2(v,u,0); 
    }
    //计算轻子结点答案 
    if(son[u]) dfs2(son[u],u,1);//重子结点 
    add(u);//添加当前结点 
    for(auto v:g[u]) {
        if(v == fa || v == son[u]) continue;
            for(int i = dfn[v]; i <= dfn[v]+sz[v]-1; i ++)
                add(rnk[i]);
    }
    //添加轻子结点 
    for(auto t:q[u]) {
        ans[t.id] = cnt[t.k];
    }
    //算答案 
    if(!f) {
        for(int i = dfn[u]; i <= dfn[u]+sz[u]-1; i ++)
            del(rnk[i]);
    }
    //若当前为轻子结点,则清空当前子树的影响 
}

int main() {
    cin >> n >> m;
    for(int i = 1; i <= n; i ++)
        cin >> c[i];
    for(int i = 1; i < n; i ++) {
        int u,v;
        cin >> u >> v;
        g[u].push_back(v);
        g[v].push_back(u);
    }
    for(int i = 1; i <= m; i ++) {
        int v,k;
        cin >> v >> k;
        q[v].push_back({i,k});
    }
    dfs1(1,0);
    dfs2(1,0,0);
    for(int i = 1; i <= m; i ++) {
        cout << ans[i] << '\n';
    }
    return 0;
}

维护子树方法 — 深度优先搜索

若题目中的信息需要用到其子结点,
则深搜遍历子树维护。

例题 CF600E Lomsat gelral

#include<iostream>
#include<vector>
#include<cstring>

using namespace std;
typedef long long ll;

const int N = 1e5+5;
int n,m;
ll c[N];
vector<int> g[N];
int sz[N],son[N];
ll ans[N];
int cnt[N],max_color;
ll now_ans;

void dfs1(int u,int fa) {
    sz[u] = 1;
    for(auto v:g[u]) {
        if(v == fa) continue;
        dfs1(v,u);
        if(sz[v] > sz[son[u]]) son[u] = v;
        sz[u] += sz[v];
    }
}

void add_now(int u) {
    cnt[c[u]] ++;
    if(cnt[c[u]] > max_color) {
        max_color = cnt[c[u]];
        now_ans = c[u];
    }
    else if(cnt[c[u]] == max_color) {
        now_ans +=  c[u];
    }
}
//添加当前点贡献

void add_subtree(int u,int fa) {
    for(auto v:g[u]) {
        if(v == fa) continue;
            add_subtree(v,u);
    }
    add_now(u);
}
//添加子树贡献

void del_subtree(int u,int fa) {
    cnt[c[u]]--;
    for(auto v:g[u]) {
        if(v == fa) continue;
        del_subtree(v,u);
    }
}
//删除贡献

void dfs2(int u,int fa,bool f) {
    for(auto v:g[u]) {
        if(v == fa || v == son[u]) continue;
        dfs2(v,u,0); 
    }
    if(son[u]) dfs2(son[u],u,1);
    for(auto v:g[u]) {
        if(v == fa || v == son[u]) continue;
            add_subtree(v,u);
    }
    add_now(u);
    ans[u] = now_ans;
    if(!f) {
        del_subtree(u,fa);
        max_color = 0;
        now_ans = 0;
    }
}

int main() {
    cin >> n;
    for(int i = 1; i <= n; i ++)
        cin >> c[i];
    for(int i = 1; i < n; i ++) {
        int u,v;
        cin >> u >> v;
        g[u].push_back(v);
        g[v].push_back(u);
    }
    dfs1(1,0);
    dfs2(1,0,0);
    for(int i = 1; i <= n; i ++) {
        cout << ans[i] << ' ';
    }
    return 0;
}

习题

CF246E Blood Cousins Return ::::info[Code]

#include<iostream>
#include<vector>
#include<unordered_map>
#include<cstring>

using namespace std;

const int N = 1e5+5;
int n,m;
vector<int> g[N],rt;
string s[N];
int dfn[N],rnk[N],sz[N],son[N],dep[N],tot = 0;
struct Query {
    int id,k;
};
vector<Query> q[N];
int ans[N];

unordered_map<string,int> vis[N];
int now_ans[N];

void dfs1(int u,int fa) {
    dfn[u] = ++tot;
    rnk[tot] = u;
    dep[u] = dep[fa]+1;
    sz[u] = 1;
    for(auto v:g[u]) {
        dfs1(v,u);
        if(sz[v] > sz[son[u]]) son[u] = v;
        sz[u] += sz[v];
    }
}

void add(int u) {
    if(!vis[dep[u]].count(s[u]))
        now_ans[dep[u]] ++;
    vis[dep[u]][s[u]] ++;
}
void del(int u) {
    vis[dep[u]][s[u]] --;
    if(vis[dep[u]][s[u]] == 0) {
        now_ans[dep[u]] --;
        vis[dep[u]].erase(s[u]);
    }
}

void dfs2(int u,bool f) {
    for(auto v:g[u]) {
        if(v == son[u]) continue;
        dfs2(v,0);
    }
    if(son[u]) dfs2(son[u],1);
    for(auto v:g[u]) {
        if(v == son[u]) continue;
            for(int i = dfn[v]; i <= dfn[v]+sz[v]-1; i ++)
                add(rnk[i]);
    }
    add(u);
    for(auto t:q[u]) {
        if(dep[u]+t.k <= n)
            ans[t.id] = now_ans[dep[u]+t.k];
    }
    if(!f) {
        for(int i = dfn[u]; i <= dfn[u]+sz[u]-1; i ++)
            del(rnk[i]);
    }
} 

int main() {
    cin >> n;
    for(int i = 1; i <= n; i ++) {
        int f;
        cin >> s[i] >> f;
        if(f == 0) rt.push_back(i);
        else g[f].push_back(i);
    }
    cin >> m;
    for(int i = 1; i <= m; i ++) {
        int v,k;
        cin >> v >> k;
        q[v].push_back({i,k});
    }
    for(auto u:rt) {
        tot = 0;
        dep[0] = 0;
        dfs1(u,0);
        dfs2(u,0);
    }

    for(int i = 1; i <= m; i ++) {
        cout << ans[i] << '\n';
    }
    return 0;
}

:::: P8844 [传智杯 #4 初赛] 小卡与落叶 ::::info[Code]

#include<iostream>
#include<vector>

using namespace std;

const int N = 1e5+5;
const int M = 2*N;
int n,m;
int head[N],nxt[M],to[M],tote = 0;
void add_edge(int u,int v) {
    nxt[++tote] = head[u]; 
    head[u] = tote;
    to[tote] = v;
}

struct BIT {
    int c[N] = {0};
    int lowbit(int x){return x&(-x);}
    void update(int u,int w)
    {
        for(int i = u; i <= n; i += lowbit(i))
            c[i] += w;
    }
    int q(int u)
    {
        int ans = 0;
        for(int i = u; i > 0; i -= lowbit(i))
            ans += c[i];
        return ans;
    }
    int query(int u) {
        return q(n)-q(u-1);
    }
}bit;

int dfn[N],rnk[N],sz[N],son[N],dep[N],tot = 0;
void dfs1(int u,int fa) {
    dfn[u] = ++tot;
    rnk[tot] = u;
    dep[u] = dep[fa]+1;
    sz[u] = 1;
    for(int i = head[u]; i; i = nxt[i]) {
        int v = to[i];
        if(v == fa) continue;
        dfs1(v,u);
        if(sz[v] > sz[son[u]]) son[u] = v;
        sz[u] += sz[v];
    }
}

struct Query {
    int id,x;
};
vector<Query> q[N];
int ans[N],totq = 0;

void add(int u) {
    bit.update(dep[u],1);
}
void del(int u) {
    bit.update(dep[u],-1);
}

void dfs2(int u,int fa,bool f) {
    for(int i = head[u]; i; i = nxt[i]) {
        int v = to[i];
        if(v == fa || v == son[u]) continue;
        dfs2(v,u,0);
    }
    if(son[u]) dfs2(son[u],u,1);
    for(int i = head[u]; i; i = nxt[i]) {
        int v = to[i];
        if(v == fa || v == son[u]) continue;
            for(int j = dfn[v]; j <= dfn[v]+sz[v]-1; j ++)
                add(rnk[j]);
    }
    add(u);
    for(auto t:q[u]) {
        ans[t.id] = bit.query(max(dep[u],t.x));
    }
    if(!f) {
        for(int i = dfn[u]; i <= dfn[u]+sz[u]-1; i ++)
            del(rnk[i]);
    }
} 

int main() {
    ios::sync_with_stdio(0);
    cin.tie(0);
    cout.tie(0);
    cin >> n >> m;
    for(int i = 1; i < n; i ++) {
        int u,v;
        cin >> u >> v;
        add_edge(u,v);
        add_edge(v,u);
    }
    int lst = n+1;
    for(int i = 1; i <= m; i ++) {
        int opt,x;
        cin >> opt >> x;
        if(opt == 1)
            lst = x;
        else
            q[x].push_back({++totq,lst});
    }
    dfs1(1,0);
    dfs2(1,0,0);
    for(int i = 1; i <= totq; i ++) {
        cout << ans[i] << '\n';
    }
    return 0;
}

::::

总结

树上启发式合并主要用于解决与子树相关的离线查询问题,
如统计子树内颜色、数值频率、众数、中位数等。

时间复杂度 O(n\log n)

典型适用问题特征

何为启发式合并

启发式

我们保留重子结点的选择,也是根据直觉进行的优化。

合并
结点 u 的答案 = 重儿子的信息 \cup 轻儿子的信息 \cup 自己的信息。

并查集的按秩合并平衡树便用了启发式合并。
通常都是将小树合并到大树上。

而在树上启发式合并中:
基于“子树越大越有价值”的经验选择,把小的信息合并到大的信息上。