DongDong數顔色 樹上啟發式合并 牛客

連結：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;
}