题目链接:旗鼓相当的对手
显然,我们可以类似树形dp,对于一个点统计答案的时候,分别从不同子树当中加入。
但是,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,a[N],res[N],sz[N],son[N],dep[N],tmp[N];
vector<int> g[N]; vector<pair<int,int>> v;
unordered_map<int,int> num,sum;
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; dep[x]=dep[fa]+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 insert(int x,int fa,int pos){
res[pos]+=num[k-dep[x]+2*dep[pos]]*a[x]+sum[k-dep[x]+2*dep[pos]];
v.push_back({dep[x],a[x]});
for(auto to:g[x]) if(to!=fa) insert(to,x,pos);
}
void del(int x,int fa){
num[dep[x]]--,sum[dep[x]]-=a[x];
for(auto to:g[x]) if(to!=fa) del(to,x);
}
void dfs2(int x,int fa){
for(auto to:g[x]) if(to!=son[x]&&to!=fa) dfs2(to,x),del(to,x);
if(son[x]) dfs2(son[x],x);
for(auto to:g[x]) if(to!=son[x]&&to!=fa){
v.clear(); insert(to,x,x);
for(auto i:v) num[i.first]++,sum[i.first]+=i.second;
}
num[dep[x]]++,sum[dep[x]]+=a[x];
}
signed main(){
cin>>n>>k;
for(int i=1;i<=n;i++) scanf("%lld",&a[i]);
for(int i=1,a,b;i<n;i++) scanf("%lld %lld",&a,&b),add(a,b);
dfs1(1,1); dfs2(1,1);
for(int i=1;i<=n;i++) printf("%lld ",res[i]);
return 0;
}