題目鏈接:旗鼓相當的對手
顯然,我們可以類似樹形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;
}