DSU on tree !
解决树上问题的利器,复杂度虽然没有长链剖分优秀,不过思考简单而且代码优美,是树上维护答案的好帮手。
例题:DSU on tree
应用范围
解决一些子树的离线静态问题,巧妙地将暴力 \(O(n^2)\) 的复杂度优化到 \(O(nlogn)\)。
算法思路
- 回溯整棵树维护子树大小以及重儿子,为后面的 DSU 做好准备。
- 对于每个结点,求出其所有轻儿子的答案,但不继承。
- 求出重儿子的答案,并继承。
- 暴力合并轻儿子答案,最后消除轻儿子对于结果的影响。
复杂度
\(O(nlogn)\)
证明
待补
CF600E
DSU on tree 的经典题,考虑询问离线后先预处理,直接做的话暴力是 \(o(n^2)\) 的。这个时候使用 DSU on tree 即可,注意维护子树内颜色的数量和暴力合并轻儿子的过程。
#include<bits/stdc++.h>
#define int long long
using namespace std;
const int N = 100000 + 10;
int n,tot,hson,sum,maxn;
int ans[N],cnt[N],col[N];
int head[N],ver[2*N],Next[2*N];
int siz[N],son[N];
void add(int x,int y)
{
ver[++tot] = y,Next[tot] = head[x],head[x] = tot;
}
void dfs(int u,int fa)
{
siz[u] = 1;
for(int i = head[u];i;i = Next[i])
{
int v = ver[i];
if(v == fa) continue;
dfs(v,u);
siz[u] += siz[v];
if(siz[v] > siz[son[u]]) son[u] = v;
}
}
void modify(int u,int fa,int val)
{
cnt[col[u]] += val;
if(cnt[col[u]] > maxn) maxn = cnt[col[u]],sum = col[u];
else if(cnt[col[u]] == maxn) sum += col[u];
for(int i = head[u];i;i = Next[i])
{
int v = ver[i];
if(v == fa || v == hson) continue;
modify(v,u,val);
}
}
void dsu(int u,int fa,int keep)
{
for(int i = head[u];i;i = Next[i])
{
int v = ver[i];
if(v == fa || v == son[u]) continue;
dsu(v,u,0);
}
if(son[u]) dsu(son[u],u,1),hson = son[u];
modify(u,fa,1); hson = 0;
ans[u] = sum;
if(not keep) modify(u,fa,-1),sum = 0,maxn = 0;
}
signed main()
{
tot = 1;
scanf("%lld",&n);
for(int i = 1;i <= n;i++) scanf("%lld",&col[i]);
for(int i = 1;i < n;i++)
{
int u,v;
scanf("%lld%lld",&u,&v);
add(u,v);
add(v,u);
}
dfs(1,0);
//cout << 1 << endl;
dsu(1,0,0);
for(int i = 1;i <= n;i++) printf("%lld ",ans[i]);
printf("\n");
return 0;
}