[HDU4358] Boring counting (树上启发式合并)

传送门
感叹一句，暴力即优雅。
CF上这个blog写得真心好：[Tutorial] Sack (dsu on tree)
我仍未知道这和合并有什么关系

#include<iostream>
#include<cstdio>
#include<cstring>
#include<map>
#define lson (node<<1)
#define rson (node<<1|1)
using namespace std;
const int N=100001;
const int M=N<<1;
int T,n,m,K,ans[N],col[N],sz[N],bc[N];
int pre[N],dfn[N],ver[N];
int cnt,Next[M],v[M],head[N];
int tot,s[N];
map<int,int> mp;

void read(int &x){
    char ch=getchar();x=0;
    for(;ch<'0'||ch>'9';ch=getchar());
    for(;ch>='0'&&ch<='9';ch=getchar()) x=(x<<3)+(x<<1)+ch-'0';
}

void add(int x,int y){
    Next[++cnt]=head[x];
    head[x]=cnt;
    v[cnt]=y;
}

void get_size(int x,int fa){
    sz[x]=1;bc[x]=0;
    pre[x]=++cnt;ver[cnt]=x;
    for(int i=head[x];i!=-1;i=Next[i])
     if (v[i]!=fa){
        get_size(v[i],x);
        if (sz[v[i]]>sz[bc[x]]) bc[x]=v[i];
        sz[x]+=sz[v[i]];
     }
    dfn[x]=cnt;
}

void color(int c,int v){
    if (s[c]==K) cnt--;
    s[c]+=v;
    if (s[c]==K) cnt++;
}

void dfs(int x,int fa,bool keep){
    for(int i=head[x];i!=-1;i=Next[i])
     if (v[i]!=fa&&v[i]!=bc[x]) dfs(v[i],x,0);
    if (bc[x]){
        dfs(bc[x],x,1);
        for(int i=pre[x]+1;i<pre[bc[x]];i++) color(col[ver[i]],1);
        for(int i=dfn[bc[x]]+1;i<=dfn[x];i++) color(col[ver[i]],1);
    }
    color(col[x],1);
    ans[x]=cnt;
    if(!keep){
        for(int i=pre[x];i<=dfn[x];i++)
         color(col[ver[i]],-1);
    }
}

int main(){
    read(T);
    for(int o=1;o<=T;o++){
        if (o>1) printf("\n");
        printf("Case #%d:\n",o);
        memset(head,-1,sizeof head);
        cnt=tot=0;mp.clear();
        read(n);read(K);
        int x,y;
        for(int i=1;i<=n;i++){
            read(x);
            if (!mp[x]) mp[x]=++tot;
            col[i]=mp[x];
        }
        for(int i=1;i<n;i++){
            read(x);read(y);
            add(x,y);add(y,x);
        }
        cnt=0;
        get_size(1,0);
        cnt=0;
        dfs(1,0,0);
        read(m);
        for(int i=1;i<=m;i++){
            read(x);
            printf("%d\n",ans[x]);
        }
    }
    return 0;
}