题目链接
题意:有一棵以 1 1 1 为根结点的树,树上每个结点有一个颜色,多组询问:问以 u u u 为根的子树上共有多少种不同的颜色。
解决该题,最暴力的做法就是对于每一棵子树都dfs一遍统计出答案,统计完清空: O ( n 2 ) O(n^2) O(n2)。但“听说”可以用dsu on tree优化:
dsu on tree [ O ( n l o g n ) O(nlogn) O(nlogn) ]
统计以结点 u u u 为根结点的子树
- 遍历结点 u u u 的轻儿子,将轻儿子的贡献清空
- 统计结点 u u u 重儿子的贡献
- 暴力统计轻儿子的贡献
#include <bits/stdc++.h>
using namespace std;
int read()
{
int x = 0, f = 1; char c = getchar();
while (c < '0' || c > '9') { if (c == '-') f = -f; c = getchar(); }
while (c >= '0' && c <= '9') { x = x * 10 + c - '0'; c = getchar(); }
return x * f;
}
const int maxN = 100005;
int n, m, c[maxN]; //c[i]: 颜色标号
struct EDGE{
int adj, to;
EDGE(int a = -1, int b = 0): adj(a), to(b) {}
}edge[maxN << 1];
int head[maxN], cnt;
void add_edge(int u, int v)
{
edge[cnt] = EDGE(head[u], v);
head[u] = cnt ++ ;
}
int siz[maxN], son[maxN]; //子树大小 //重儿子标号
void dfs(int u, int fa)
{
siz[u] = 1;
for(int i = head[u]; ~i; i = edge[i].adj)
{
int v = edge[i].to;
if(v == fa) continue;
dfs(v, u);
siz[u] += siz[v];
if(siz[son[u]] < siz[v]) son[u] = v;
}
}
int color[maxN], ans[maxN], tot, nowSon;
/*
* color[i]: 颜色标号为i的颜色出现的次数
* ans[i]: 以i节点为根节点的子树上的颜色总数
* tot: 当前不同颜色的个数
* nowSon: 当前的重儿子
*/
void cal(int u, int fa, int val)
{
if(!color[c[u]]) ++ tot;
color[c[u]] += val;
for(int i = head[u]; ~i; i = edge[i].adj)
{
int v = edge[i].to;
if(v == fa || v == nowSon) continue;
cal(v, u, val);
}
}
void dsu(int u, int fa, bool op) //op==0: 轻边,需要清空
{
for(int i = head[u]; ~i; i = edge[i].adj) //遍历轻儿子
{
int v = edge[i].to;
if(v == fa || v == son[u]) continue;
dsu(v, u, false);
}
if(son[u]) dsu(son[u], u, true), nowSon = son[u]; //统计重儿子的贡献
cal(u, fa, 1), nowSon = 0; // 暴力统计轻儿子的贡献(记录重儿子nowSon来忽略重儿子的贡献)
ans[u] = tot; //得到u子树的答案
if(!op) { cal(u, fa, -1); tot = 0; } // 如果是轻儿子,清除该子树的影响。
}
int main()
{
memset(head, -1, sizeof(head));
n = read();
for(int i = 0; i < n - 1; ++ i )
{
int u, v; u = read(); v = read();
add_edge(u, v); add_edge(v, u);
}
for(int i = 1; i <= n; ++i ) c[i] = read();
dfs(1, 0);
dsu(1, 0, true);
m = read();
for(int i = 0; i < m; ++ i )
{
int u = read();
printf("%d\n", ans[u]);
}
return 0;
}
}
参考博客:https://www.cnblogs.com/TY02/p/11218999.html【实践证明,调试一遍就会了】
我的调试样例:
7
1 2
1 3
2 4
2 5
4 6
3 7
1 2 2 3 3 4 3
7
1
2
3
4
5
6
7