[USACO12FEB]附近的牛Nearby Cows

洛咕

题意:给你一棵\(n\)个点的树，点带权，对于每个节点求出距离它不超过 \(k\)的所有节点权值和\(a_i\).

分析:"对于每个节点"即相当于要以每个节点为根算一次贡献,还是考虑换根\(DP\).设\(f[i][j]\)表示以\(i\)点为根的子树内与\(i\)距离不超过\(j\)的所有节点的点权和,则\(f[u][j]+=f[v][j-1],u->v\).

那么以\(1\)号点为根求出上述\(f\)数组后,\(ans[1]=f[1][k]\).那么对于\(u->v\),如何由\(ans[u]\)求出\(ans[v]\)呢?随便手玩一棵树发现转移的时候,需要知道\(g[i][k]\)表示以\(i\)点为中心,与\(i\)距离不超过\(k\)的所有点的点权和.这个可以在换根时通过\(f[i][j]\)得到.写到这里我突然发现我的\(ans\)数组是多余的啊,换根\(DP\)的时候我都已经把所有的\(g[i][j]\)都求出来了,最终答案不就是每个点的\(g[i][k]\)么?

#include<iostream>
#include<cstdio>
#include<algorithm>
#include<cstring>
#include<cmath>
#include<queue>
#include<map>
#include<set>
#define ll long long
using namespace std;
inline int read(){
    int x=0,o=1;char ch=getchar();
    while(ch!='-'&&(ch<'0'||ch>'9'))ch=getchar();
    if(ch=='-')o=-1,ch=getchar();
    while(ch>='0'&&ch<='9')x=x*10+ch-'0',ch=getchar();
    return x*o;
}
const int N=1e5+5;
int n,k,a[N];ll g[N][21],f[N][21],ans[N];
int tot,head[N],nxt[N<<1],to[N<<1];
inline void add(int a,int b){
	nxt[++tot]=head[a];head[a]=tot;to[tot]=b;
}
inline void dfs(int u,int fa){
	f[u][0]=a[u];
	for(int i=head[u];i;i=nxt[i]){
		int v=to[i];if(v==fa)continue;
		dfs(v,u);
		for(int j=1;j<=k;++j){
			f[u][j]+=f[v][j-1];
		}
	}
}
inline void dp(int u,int fa){
	for(int i=head[u];i;i=nxt[i]){
		int v=to[i];if(v==fa)continue;
		ans[v]=ans[u]-g[u][k]+f[v][k-1]+f[v][k];
		for(int j=1;j<=k;++j){
			g[v][j]=f[v][j]+g[u][j-1];
			if(j>=2)g[v][j]-=f[v][j-2];
		}
		dp(v,u);
	}
}
int main(){
	n=read();k=read();
	for(int i=1;i<n;++i){
		int a=read(),b=read();
		add(a,b);add(b,a);
	}
	for(int i=1;i<=n;++i)a[i]=read();
	dfs(1,0);
	for(int j=0;j<=k;++j)ans[1]+=f[1][j],g[1][j]=f[1][j];
	for(int i=1;i<=n;++i)g[i][0]=f[i][0];
	dp(1,0);
	for(int i=1;i<=n;++i)printf("%lld\n",ans[i]);
    return 0;
}