HDU-4358：Boring counting（树上启发式合并）

Boring counting
Time Limit: 6000/3000 MS (Java/Others)
Memory Limit: 98304/98304 K (Java/Others)

Problem Description
In this problem we consider a rooted tree with N vertices. The vertices are numbered from 1 to N, and vertex 1 represents the root. There are integer weights on each vectice. Your task is to answer a list of queries, for each query, please tell us among all the vertices in the subtree rooted at vertice u, how many different kinds of weights appear exactly K times?

Input
The first line of the input contains an integer T(T<=5) T ( T <= 5 ) , indicating the number of test cases.
For each test case, the first line contains two integers N and K, as described above. (1<=N<=105,1<=K<=N) ( 1 <= N <= 10 5 , 1 <= K <= N )
Then come N integers in the second line, they are the weights of vertice 1 to N. (0<=weight<=109) ( 0 <= w e i g h t <= 10 9 )
For next N-1 lines, each line contains two vertices u and v, which is connected in the tree.
Next line is a integer Q, representing the number of queries. (1<=Q<=105) ( 1 <= Q <= 10 5 )
For next Q lines, each with an integer u, as the root of the subtree described above.

Output
For each test case, output “Case #X:” first, X is the test number. Then output Q lines, each with a number – the answer to each query.

Seperate each test case with an empty line.

Sample Input
1
3 1
1 2 2
1 2
1 3
3
2
1
3

Sample Output
Case #1:
1
1
1

思路：树上启发式合并貌似很好做？

#include
using namespace std;
const int MAX=1e5+10;
vector<int>e[MAX];
vector<int>v;
int a[MAX],siz[MAX],L[MAX],R[MAX],son[MAX],num[MAX],all;
void dfs(int k,int fa)
{
    siz[k]=1;
    L[k]=++all;
    num[all]=a[k];
    son[k]=0;
    for(int i=0;iint nex=e[k][i];
        if(nex==fa)continue;
        dfs(nex,k);
        siz[k]+=siz[nex];
        if(siz[son[k]]int ans[MAX];
int cnt[MAX];
int tot;
int n,m;
void cal(int k,int fa,int tp)//tp=0/1表示k不是/是重儿子
{
    for(int i=0;iint nex=e[k][i];
        if(nex==fa||nex==son[k])continue;//优先遍历轻儿子
        cal(nex,k,0);
    }
    if(son[k])cal(son[k],k,1);//最后遍历重儿子，此时重儿子的信息已被存储下来
    cnt[a[k]]++;
    if(cnt[a[k]]==m)tot++;
    if(cnt[a[k]]==m+1)tot--;
    for(int i=0;i//暴力统计轻儿子对答案的影响
    {
        int nex=e[k][i];
        if(nex==fa||nex==son[k])continue;
        for(int j=L[nex];j<=R[nex];j++)
        {
            cnt[num[j]]++;
            if(cnt[num[j]]==m)tot++;
            if(cnt[num[j]]==m+1)tot--;
        }
    }
    ans[k]=tot;
    if(tp==0)//如果当前节点不是重儿子，清除信息
    {
        tot=0;
        for(int j=L[k];j<=R[k];j++)cnt[num[j]]--;
    }
}
int main()
{
    int T;
    cin>>T;
    for(int cas=1;cas<=T;cas++)
    {
        scanf("%d%d",&n,&m);
        v.clear();
        for(int i=1;i<=n;i++)
        {
            scanf("%d",&a[i]);
            v.push_back(a[i]);
        }
        sort(v.begin(),v.end());
        v.erase(unique(v.begin(),v.end()),v.end());//将点的权值离散化
        for(int i=1;i<=n;i++)a[i]=lower_bound(v.begin(),v.end(),a[i])-v.begin()+1;
        for(int i=1;i<=n;i++)e[i].clear();
        for(int i=1;i<=n;i++)cnt[i]=0;
        for(int i=1;iint x,y;
            scanf("%d%d",&x,&y);
            e[x].push_back(y);
            e[y].push_back(x);
        }
        all=0;
        dfs(1,0);
        tot=0;
        cal(1,0,1);
        int Q;
        scanf("%d",&Q);
        printf("Case #%d:\n",cas);
        while(Q--)
        {
            int x;
            scanf("%d",&x);
            printf("%d\n",ans[x]);
        }
        if(casprintf("\n");
    }
    return 0;
}