luogu P3047 [USACO12FEB]Nearby Cows G
题面传送门
显然可以\(dp\)
第一次\(dfs\),定义\(f_{i,j}\)为在\(i\)子树下距离在\(j\)以内的总和。
第二次可以容斥,用父亲节点的减去儿子节点的值。
代码实现:
#include<cstdio>
#include<cstring>
using namespace std;
int n,m,k,f[100039][39],dp[100039][39],a[100039],x,y,h[100039],head;
struct yyy {
int to,z;
} s[200039];
inline void add(int x,int y) {
s[++head]=(yyy) {
y,h[x]
};
h[x]=head;
}
inline void dfs1(int x,int last) {
int cur=h[x],i;
yyy tmp;
for(i=0; i<=m; i++) f[x][i]=a[x];
while(cur!=-1) {
tmp=s[cur];
if(tmp.to!=last) {
dfs1(tmp.to,x);
for(i=1; i<=m; i++) f[x][i]+=f[tmp.to][i-1];
}
cur=tmp.z;
}
return;
}
inline void dfs2(int x,int last) {
int cur=h[x],i;
yyy tmp;
if(x==1)for(i=0; i<=m; i++) dp[x][i]=f[x][i];
else {
dp[x][0]=a[x];dp[x][1]=f[x][1]+dp[last][0];
for(i=2; i<=m; i++) dp[x][i]=dp[last][i-1]-f[x][i-2]+f[x][i];
}
while(cur!=-1){
tmp=s[cur];
if(tmp.to!=last) dfs2(tmp.to,x);
cur=tmp.z;
}
}
int main() {
memset(h,-1,sizeof(h));
register int i;
scanf("%d%d",&n,&m);
for(i=1; i<n; i++)scanf("%d%d",&x,&y),add(x,y),add(y,x);
for(i=1; i<=n; i++)scanf("%d",&a[i]);
dfs1(1,0);
dfs2(1,0);
for(i=1;i<=n;i++) printf("%d\n",dp[i][m]);
}