LUOGU P3047 [USACO12FEB]附近的牛Nearby Cows

传送门

解题思路

树形dp，看到数据范围应该能想到是O(nk)级别的算法，进而就可以设出dp状态，dp[x][j]表示以x为根的子树，距离它为i的点的总和，第一遍dp首先自底向上，dp出每个节点的子树中到他距离为j的，转移方程dp[x][j]=dp[u][j-1] ，第二遍dp自顶向下，dp出每个节点父亲那头的距离为j的，转移方程dp[u][j]+=dp[x][j-1]-dp[u][j-2]

代码

//dp[x][j] 以i为根的子树，距离为j的牛的和
//dp[x][j]+=dp[u][j-1];
//dp[u][j]+=dp[x][j-1]-dp[u][j-2];
#include<iostream>
#include<cstdio>
#include<cstring>
#include<algorithm>

using namespace std;
const int MAXN = 100005;

inline int rd(){
	int x=0,f=1;char ch=getchar();
	while(!isdigit(ch)) {f=ch=='-'?0:1;ch=getchar();}
	while(isdigit(ch))  {x=(x<<1)+(x<<3)+ch-'0';ch=getchar();}
	return f?x:-x;
}

int dp[MAXN][25],head[MAXN],cnt;
int n,k,ans[MAXN],to[MAXN<<1],nxt[MAXN<<1];

inline void add(int bg,int ed){
	to[++cnt]=ed,nxt[cnt]=head[bg],head[bg]=cnt;
}

void dfs1(int x,int fa){
	for(register int i=head[x];i;i=nxt[i]){
		int u=to[i];if(u==fa) continue;
		dfs1(u,x);
		for(register int j=1;j<=k;j++) dp[x][j]+=dp[u][j-1];
	}
}

void dfs2(int x,int fa){
	for(register int i=head[x];i;i=nxt[i]){
		int u=to[i];if(u==fa) continue;
		for(register int j=k;j>=2;j--)
			dp[u][j]=dp[u][j]+dp[x][j-1]-dp[u][j-2];
		dp[u][1]+=dp[x][0];
		dfs2(u,x);
	}
}

int main() {
	n=rd(),k=rd();int x,y;
	for(int i=1;i<n;i++){
		x=rd(),y=rd();
		add(x,y),add(y,x);
	}
	for(int i=1;i<=n;i++) dp[i][0]=rd();
	dfs1(1,0);dfs2(1,0);
	for(int i=1;i<=n;i++) {
		int ans=0;
		for(register int j=0;j<=k;j++)
			ans+=dp[i][j];
		printf("%d\n",ans);
	}
	return 0;
}