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;
    }
}