旗鼓相當的對手

題目鏈接:旗鼓相當的對手


顯然,我們可以類似樹形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;
}
發表評論
所有評論
還沒有人評論,想成為第一個評論的人麼? 請在上方評論欄輸入並且點擊發布.
相關文章