题:https://ac.nowcoder.com/acm/contest/4853/E
题意:对于一对点(u,v)要是dis(u,v)==k,就会对这对点的LCA产生a[x]+a[y]的贡献(LCA!=u&&LCA!=v)
分析:简单的dsu,把每个点当作LCA去统计子树深度个数,经过LCA的路径就是深度之和;
对于每个点u深度对应要达成k的深度要为2*deep[LCA]-deep[u]+k;
#include<bits/stdc++.h> using namespace std; #define pb push_back typedef long long ll; const int M=2e5+5; ll ans[M]; ll val[M],sz[M],son[M],a[M],cnt[M],deep[M]; ll k,n; vector<ll>g[M]; ll nownum; void dfs1(ll u,ll fa){ sz[u]=1; for(auto v:g[u]){ if(v!=fa){ deep[v]=deep[u]+1; dfs1(v,u); sz[u]+=sz[v]; if(sz[v]>sz[son[u]]) son[u]=v; } } } void add(ll u,ll fa,ll x){ cnt[deep[u]]+=x; val[deep[u]]+=x*a[u]; for(auto v:g[u]){ if(v!=fa) add(v,u,x); } } void cal(ll u,ll fa,ll LCA){ ll need=k+deep[LCA]*2-deep[u]; if(need>0){ nownum+=val[need]; nownum+=cnt[need]*a[u]; } for(auto v:g[u]) if(v!=fa) cal(v,u,LCA); } void dfs2(ll u,ll fa,ll sign){ //cout<<u<<endl; for(auto v:g[u]){ if(v!=fa&&son[u]!=v) dfs2(v,u,0); } if(son[u]) dfs2(son[u],u,1); for(auto v:g[u]) if(v!=fa&&v!=son[u]) cal(v,u,u),add(v,u,1); ans[u]=nownum; cnt[deep[u]]++,val[deep[u]]+=a[u]; if(!sign) add(u,fa,-1); nownum=0; } int main(){ scanf("%lld%lld",&n,&k); for(ll i=1;i<=n;i++) scanf("%lld",&a[i]); for(ll u,v,i=1;i<n;i++){ scanf("%lld%lld",&u,&v); g[u].pb(v); g[v].pb(u); } dfs1(1,0); dfs2(1,0,1); for(ll i=1;i<=n;i++) printf("%lld ",ans[i]); return 0; }