2012 Multi-University Training Contest 6-1009 hdu4358 Boring counting

http://acm.hdu.edu.cn/showproblem.php?pid=4366

树形启发式合并，利用map

启发式合并，不管父子关系如何，每次将元素个数少的节点并入元素个数多的。用id[i]表示节点i合并后的结果存在节点id[i]。

#include <iostream>
#include <cstdio>
#include <memory.h>
#include <algorithm>
#include <cmath>
#include <string>
#include <climits>
#include <vector>
#include <set>
#include <utility>
#include <map>
#pragma comment(linker, "/STACK:102400000,102400000")
using namespace std;
#define N 100010
map<int,int> m[N];
vector<int> emp[N],emp0[N];
int id[N],w[N],ans[N];
bool vis[N];
int n,k;
void build(int rt)
{
    vis[rt]=true;
    for (int i=0;i<(int)emp0[rt].size();i++)
    {
        int son=emp0[rt][i];
        if (!vis[son])
        {
            emp[rt].push_back(son);
            build(son);
        }
    }
}
void merge(int rt,int src,int dest)
{
    map<int,int>::iterator itr,itr1;
    for (itr=m[src].begin();itr!=m[src].end();itr++)
    {
        itr1=m[dest].find(itr->first);
        if (itr1!=m[dest].end())
        {
            if (itr1->second==k)
            {
                ans[rt]--;
                itr1->second+=itr->second;
            }
            else
            {
                itr1->second+=itr->second;
                if (itr1->second==k)
                    ans[rt]++;
            }
        }
        else
        {
            m[dest][itr->first]=itr->second;
            if (itr->second==k)
               ans[rt]++;
        }
    }
}
void dfs(int rt)
{
    int l=emp[rt].size();
    if (l==0) return;

    for (int i=0;i<l;i++)
    {
        int son=emp[rt][i];
        dfs(son);
        int l1=emp[id[rt]].size();
        int l2=emp[id[son]].size();
        if (l1>l2)
        {
            merge(rt,id[son],id[rt]);
        }
        else
        {
            ans[rt]=ans[son];
            merge(rt,id[rt],id[son]);
            id[rt]=id[son];
        }
    }
}
int main()
{
    int t,x,y;
    scanf("%d",&t);
    for (int cas=1;cas<=t;cas++)
    {
        scanf("%d%d",&n,&k);
        for (int i=1;i<=n;i++)
            scanf("%d",&w[i]);


        for (int i=0;i<N;i++)
        {
            emp0[i].clear();
            emp[i].clear();
            m[i].clear();
        }
        for (int i=0;i<n-1;i++)
        {
            scanf("%d%d",&x,&y);
            emp0[x].push_back(y);
            emp0[y].push_back(x);
        }
        memset(vis,false,sizeof(vis));
        build(1);
        memset(ans,0,sizeof(ans));
        for (int i=1;i<=n;i++)
        {
            id[i]=i;
            m[i][w[i]]=1;
            if (k==1) ans[i]=1;
        }

        dfs(1);

        printf("Case #%d:\n",cas);
        int q;
        scanf("%d",&q);
        while (q--)
        {
            scanf("%d",&x);
            printf("%d\n",ans[x]);
        }
        if (cas!=t) printf("\n");
    }
    return 0;
}