P3047 [USACO12FEB]Nearby Cows G

题解

题意：给你一棵 n 个点的树，点带权，对于每个节点求出距离它不超过 k 的所有节点权值和 m​。

数据范围 n<=1e5,k<=20;

思路：考虑换根dp的转移，加上容斥，设f[i][j]表示与i这个节点相距j的节点个数，g[i][j]表示子树内与i节点相差j的节点个数；
假设已经知道父节点的一些信息，求子节点的话，发现需要加上f[fa][j-1]-g[i][j-2]+g[i][j];
g数组不用转移，就ok了

#include<bits/stdc++.h>
using namespace std;
const int N=1e5+7;
struct edge{
	int v,nxt;
}e[N<<1];
int n,cnt,k;
int a[N],head[N],fa[N];
int f[N][30],d[N][30];
void add_edge(int u,int v){
	e[++cnt]=(edge){v,head[u]};
	head[u]=cnt;
}

void build(int u){
	for(int i=head[u];i;i=e[i].nxt){
		int to=e[i].v;
		if(to!=fa[u]){
			fa[to]=u;
			build(to);
			for(int j=1;j<=k;j++){
				d[u][j]+=d[to][j-1];
			}
		}
	}
}

void dfs(int u){
	if(u==1){
		for(int i=1;i<=k;i++){
			f[u][i]=d[u][i];
		}
		for(int i=head[u];i;i=e[i].nxt){
			int to=e[i].v;
			if(to!=fa[u]){
				dfs(to);
			}
		}
	}else{
		for(int i=1;i<=k;i++){
			f[u][i]=f[fa[u]][i-1]-d[u][i-2]+d[u][i];
		}
		for(int i=head[u];i;i=e[i].nxt){
			int to=e[i].v;
			if(to!=fa[u]){
				dfs(to);
			}
		}
	}
}

int main(){
	scanf("%d%d",&n,&k);
	for(int i=1;i<n;i++){
		int x,y;
		scanf("%d%d",&x,&y);
		add_edge(x,y);
		add_edge(y,x);
	}
	for(int i=1;i<=n;i++){
		scanf("%d",&a[i]);
		f[i][0]=d[i][0]=a[i];
		for(int j=1;j<=k;j++){
			d[i][j]=a[i];
		}
	}
	build(1);
//	for(int i=1;i<=n;i++){
//		for(int j=0;j<=k;j++){
//			cout<<d[i][j]<<" ";
//		}
//		cout<<"
";
//	}
	dfs(1);
	for(int i=1;i<=n;i++) cout<<f[i][k]<<"
";
}