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

#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
;
vector<int> emp
,emp0
;
int id
,w
,ans
;
bool vis
;
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;
}