題目鏈接:HDU - 4358
直接維護當前出現次數爲k的個數,存重兒子答案,暴力計算輕兒子即可。
AC代碼:
#pragma GCC optimize("-Ofast","-funroll-all-loops")
#include<bits/stdc++.h>
//#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\n",res[x]);
}
signed main(){
int T; cin>>T;
for(int i=1;i<=T;i++){
printf("Case #%d:\n",i),solve();
if(i!=T) puts("");
}
return 0;
}