U41492 树上数颜色 题解

U41492 树上数颜色

题目描述

给一棵根为1的树，每次询问子树颜色种类数

输入格式

第一行一个整数n，表示树的结点数

接下来n-1行，每行一条边

接下来一行n个数，表示每个结点的颜色c[i]

接下来一个数m，表示询问数

接下来m行表示询问的子树

输出格式

对于每个询问，输出该子树颜色数

输入输出样例

输入 #1复制

5

1 2

1 3

2 4

2 5

1 2 2 3 3

5

1

2

3

4

5

输出 #1复制

3

2

1

1

1

说明/提示

对于前三组数据，有1≤m,c[i]≤n≤100

而对于所有数据，有1≤m,c[i]≤n≤1e5

暴力算法 ac

?

1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51

#include<iostream> #include<vector> using namespace std; const int maxn=1e5+10; vector<int> tr[maxn]; int tong[maxn],ans[maxn],c[maxn],fa[maxn],now; void calc(int u,int f,int fl) {     tong[c[u]]+=fl;     if (fl==1&&tong[c[u]]==1) now++;     if (fl==-1&&tong[c[u]]==0) now--;     for (int i=0;i<tr[u].size();i++)     {         int v=tr[u][i];         if (v==f) continue;         calc(v,u,fl);     } } void dfs(int u,int f) {     for (int i=0;i<tr[u].size();i++)     {         int v=tr[u][i];         if (v==f) continue;         dfs(v,u);     }     calc(u,f,1);     ans[u]=now;     calc(u,f,-1); } int main() {     int n,m;     cin>>n;     for (int i=1;i<n;i++)     {         int x,y;         cin>>x>>y;         tr[x].push_back(y);         tr[y].push_back(x);     }     for (int i=1;i<=n;i++) cin>>c[i];     dfs(1,0);     cin>>m;     while(m--)     {         int x;         cin>>x;         cout<<ans[x]<<endl;      } }