牛客练习赛60E 旗鼓相当的对手

dsu on tree

题目链接

点我跳转

题目大意

给你一棵以\(1\)为根节点，包含\(n\)个节点的树和一个参数 \(k\)，求每个节点的"\(rating\)"

\(rating\) 值的计算方式是这样的,对于\(u\)的子树中的所有节点,如果\(x,y\)满足\(dis(x,y) = k\)

并且\(x,y\)的最近公共祖先是\(u\)且满足\(u != x , u != y\),那么\(u\)的\(rating\)就会增加\(a_x + a_y\)

解题思路

因为 \(x , y\) 的最近公共祖先为 \(u\) ，所以 \(x , y\) 一定在 \(u\) 子树的不同分支

\(dis(x,y) = k\) 等价于 \(dep[x] + dep[y] - 2 × dep[lca(x,y)] = k\)

于是可以先用 cnt[dep] 记录深度为 dep 的节点出现的次数，用 sum[dep] 记录 dep 的节点的权值和

那么对于 \(rt\) 为根的子节点 \(u\)，与其相匹配的点的深度为 \(d = k + 2 * dep[rt] - dep[u]\)

它对 \(rt\) 产生的贡献就为 \(cnt_d × a_u\) ，而深度为 \(d\) 的点因为 \(u\) 的出现对 \(rt\) 的贡献都会翻倍

所以 \(u\) 节点的出现对 \(rt\) 的总贡献为 \(sum[d] + a[u] * cnt[d]\)

又因为与 \(u\) 节点产生贡献的节点必须和 \(u\) 不在一个分支，即一个分支内的任意节点不能相互影响

所以需要先对一个分支统计完贡献后，再添加它的信息

AC_Code

#include<bits/stdc++.h>
#define rep(i,a,n) for (int i=a;i<=n;i++)
#define int long long
using namespace std;
const int N = 3e5 + 10;
struct Edge{
	int nex , to;
}edge[N << 1];
int head[N] , TOT;
void add_edge(int u , int v)
{
	edge[++ TOT].nex = head[u] ;
	edge[TOT].to = v;
	head[u] = TOT;
}
int n , k , sum[N] , cnt[N] , ans[N];
int dep[N] , sz[N] , HH , hson[N] , f[N][30] , a[N];
void dfs(int u , int far)
{
	sz[u] = 1;
	dep[u] = dep[far] + 1;
	for(int i = head[u] ; i ; i = edge[i].nex)
	{
		int v = edge[i].to;
		if(v == far) continue ;
		dfs(v , u);
		sz[u] += sz[v];
		if(sz[v] > sz[hson[u]]) hson[u] = v;
	}
}
void change(int u , int far , int val)
{
	sum[dep[u]] += val * a[u];
	cnt[dep[u]] += val;
	for(int i = head[u] ; i ; i = edge[i].nex)
	{
		int v = edge[i].to;
		if(v == far || v == HH) continue ;
		change(v , u , val);
	}
} 
void calc(int u , int far , int rt)
{
	int c = k + 2 * dep[rt] - dep[u];
	if(c < 0) return ;
	ans[rt] += sum[c] + a[u] * cnt[c];
	for(int i = head[u] ; i ; i = edge[i].nex)
	{
		int v = edge[i].to;
		if(v == far || v == HH) continue ;
		calc(v , u , rt);
	}
}
void dsu(int u , int far , int op)
{
	for(int i = head[u] ; i ; i = edge[i].nex)
	{
		int v = edge[i].to;
		if(v == far || v == hson[u]) continue ;
		dsu(v , u , 0);
	}
	if(hson[u]) dsu(hson[u] , u , 1) , HH = hson[u];
	for(int i = head[u] ; i ; i = edge[i].nex)
	{
		int v = edge[i].to;
		if(v == far || v == HH) continue ;
		calc(v , u , u) , change(v , u , 1);
	}
	HH = 0;
	sum[dep[u]] += a[u] , cnt[dep[u]] ++ ;
	if(!op) 
	{
		change(u , far , -1);
	}
}
signed main()
{
	cin >> n >> k;
	rep(i , 1 , n) cin >> a[i]; 
	rep(i , 2 , n)
	{
		int u , v;
		cin >> u >> v;
		add_edge(u , v) , add_edge(v , u);
	}
	dfs(1 , 0);
	dsu(1 , 0 , 0);
	rep(i , 1 , n) cout << ans[i] << " \n"[i == n];
	return 0;
}