HDU - 4358

제목 링크: HDU - 4358

현재 출현 횟수가 k인 개수를 직접 유지하고 아들의 답을 중시하며 폭력은 아들을 가볍게 계산하면 된다.
AC 코드:

#pragma GCC optimize("-Ofast","-funroll-all-loops")
#include
//#define int long long
using namespace std;
const int N=1e5+10;
int n,k,q,sz[N],w[N],son[N],res[N],ans,mp[N];
vector<int> g[N],v;
inline void add(int a,int b){g[a].push_back(b),g[b].push_back(a);}
void dfs1(int x,int fa){
	sz[x]=1;
	for(auto to:g[x])	if(to!=fa){
		dfs1(to,x);	sz[x]+=sz[to];
		if(sz[to]>sz[son[x]])	son[x]=to;
	}
}
void del(int x,int fa){
	if(mp[w[x]]==k)	ans--;	mp[w[x]]--;
	for(auto to:g[x])	if(to!=fa)	del(to,x);
}
void insert(int x,int fa){
	if(mp[w[x]]==k)	ans--;	if(++mp[w[x]]==k)	ans++;
	for(auto to:g[x])	if(to!=fa)	insert(to,x);
}
void dfs2(int x,int fa){
	for(auto to:g[x])	if(to!=fa&&to!=son[x]){
		dfs2(to,x);	del(to,x);
	}
	if(son[x])	dfs2(son[x],x);
	for(auto to:g[x])	if(to!=son[x]&&to!=fa)	insert(to,x);
	if(mp[w[x]]==k)	ans--;	if(++mp[w[x]]==k)	ans++;	
	res[x]=ans;
}
void init(){
	for(int i=1;i<=n;i++)	g[i].clear(),son[i]=mp[i]=0;
	ans=0;	v.clear();
}
inline void solve(){
	cin>>n>>k;	init();
	for(int i=1;i<=n;i++)	scanf("%d",&w[i]),v.push_back(w[i]);
	sort(v.begin(),v.end()),v.erase(unique(v.begin(),v.end()),v.end());
	for(int i=1;i<=n;i++)	w[i]=lower_bound(v.begin(),v.end(),w[i])-v.begin()+1;
	for(int i=1,a,b;i<n;i++)	scanf("%d %d",&a,&b),add(a,b);
	dfs1(1,1);	dfs2(1,1);
	cin>>q;
	for(int i=1,x;i<=q;i++)	scanf("%d",&x),printf("%d
",res[x]);
}
signed main(){
	int T;	cin>>T;
	for(int i=1;i<=T;i++){
		printf("Case #%d:
",i),solve();
		if(i!=T)	puts("");
	}
	return 0;
}