链接:https://ac.nowcoder.com/acm/contest/904/E
题面:
DongDong是个喜欢数颜色的女孩子,她已经熟练地掌握了在序列上数颜色的操作,现在她开始学习如何在树上数颜色,现在给定一个n个点,n-1条边的树形图(视1号店为根),每个点有一个颜色,每次询问以x为根的子树中有多少种不同的颜色,DongDong轻松地解决了这个问题,但她想考考会编程的你。
思路:
两种做法:
1)树上启发式合并 O(nlogn) 直接套模板即可
2) 树上莫队 O(nsqrt(n))
#include <stdio.h>
#include <vector>
#include<bits/stdc++.h>
using namespace std;
const int _N = 100005;
vector<int> G[_N];
int hvs[_N], siz[_N], cnt[_N], A[_N], ans[_N], mx;
bool vis[_N];
void connect(int p, int dad)
{
siz[p] = 1;
for (int i = G[p].size() - 1; i >= 0; --i) {
int v = G[p][i];
if (v == dad) continue;
connect(v, p);
siz[p] += siz[v];
if (!hvs[p] || siz[hvs[p]] < siz[v]) hvs[p] = v;
}
return;
}
void clear(int p, int dad)
{
if(vis[A[p]]){
mx--;
vis[A[p]]=0;
}
for (int i = G[p].size() - 1; i >= 0; --i) {
int v = G[p][i];
if (v == dad) continue;
clear(v, p);
}
return;
}
void insert(int p, int dad)
{
if(!vis[A[p]]){
mx++;
vis[A[p]]=1;
}
for (int i = G[p].size() - 1; i >= 0; --i) {
int v = G[p][i];
if (v == dad) continue;
insert(v, p);
}
return;
}
void dfs(int p, int dad)
{
for (int i = G[p].size() - 1; i >= 0; --i) {
int v = G[p][i];
if (v == dad || v == hvs[p]) continue;
dfs(v, p);
clear(v, p);
mx = 0;
}
if (hvs[p]) dfs(hvs[p], p);
for (int i = G[p].size() - 1; i >= 0; --i) {
int v = G[p][i];
if (v == dad || v == hvs[p]) continue;
insert(v, p);
}
ans[p] =mx+=vis[A[p]]?0:1;
vis[A[p]]=1;
return;
}
int main()
{
int N,M;
scanf("%d%d",&N,&M);
for (int i = 1; i <= N; ++i)
scanf("%d", &A[i]);
for (int a, b, i = 1 ; i < N; ++i) {
scanf("%d%d", &a, &b);
G[a].push_back(b), G[b].push_back(a);
}
connect(1, 0);
dfs(1, 0);
for (int i = 1; i <= M; ++i){
int x;
scanf("%d",&x);
printf("%d\n", ans[x]);
}
return 0;
}