树上数颜色

7899：树上数颜色
时间限制(普通/Java):3000MS/9000MS 内存限制:250000KByte

描述

给定一个n个点、n-1条边的树（编号为1~n，其中1号为根），每个点有一个颜色，每次询问以x为根的子树中有多少种不同的颜色。

输入

第一行两个整数n，m（2<=n<=100000， 1<=m<=n），表示有n个顶点，m次查询。
第二行n个整数，表示每个点的颜色（颜色值在1~n之间）。
接下来n-1行每行u和v，表示存在一条从u到v的双向边。
接下来有m行，每行一个正整数x，表示待查询的子树的根（1<=u, v, x<=n）。
数据保证是一棵树。

输出

输出m行，每行为查询的结果。

样例输入

4 3
1 1 2 3
1 2
2 3
1 4
1
2
4

样例输出

3
2
1

思路

暴力深搜，深搜到叶子结点后回溯从下往上把集合往上合并，记录以每个根节点的子树的颜色个数

AC代码

#include <bits/stdc++.h>
using namespace std;
int n,m;
vector<int>bian[100010];
int color[100010];
set<int> dfs(int k,int fa)
{
    set<int>temp{color[k]};
    for(auto i:bian[k])
    {
        if(i!=fa)
        {
            set<int>tem=dfs(i,k);
            temp.insert(tem.begin(),tem.end());
        }
    }
    color[k]=temp.size();
    return temp;
}
int main()
{
    ios_base::sync_with_stdio(false);
    cin.tie(nullptr);
    cout.tie(nullptr);
    cin>>n>>m;
    int st,en;
    for(int i=1;i<=n;i++)
    {
        cin>>color[i];
    }
    for(int i=1;i<n;i++)
    {
        cin>>st>>en;
        bian[st].push_back(en);
        bian[en].push_back(st);
    }
    dfs(1,1);
    while (m--)
    {
        cin>>st;
        cout<<color[st]<<endl;
    }
}