P3047 [USACO12FEB]Nearby Cows G 题解

 1 #include<iostream>
 2 #include<cstring>
 3 #define N 100010
 4 using namespace std;
 5 int n,k,u1,v1;
 6 int c[N],ans[N],f[N][25];
 7 struct EDGE{
 8     int v,nxt;
 9 }edge[N*2];
10 int head[N],cnt;
11 void add(int u,int v)
12 {
13     edge[++cnt].v=v;
14     edge[cnt].nxt=head[u];
15     head[u]=cnt;
16 }
17 void dfs1(int now,int ff)
18 {
19     f[now][0]+=c[now];
20     for(int i=head[now];i!=-1;i=edge[i].nxt)
21         if(edge[i].v!=ff)
22         {
23             dfs1(edge[i].v,now);
24             for(int j=1;j<=k;j++)
25                 f[now][j]+=f[edge[i].v][j-1];
26         }
27 }
28 void dfs2(int now,int ff)
29 {
30     for(int i=head[now];i!=-1;i=edge[i].nxt)
31         if(edge[i].v!=ff)
32         {
33             for(int j=k;j>=2;j--)
34                 f[edge[i].v][j]-=f[edge[i].v][j-2];
35             for(int j=2;j<=k;j++)
36                 f[edge[i].v][j]+=f[now][j-1];
37             f[edge[i].v][1]+=c[now];
38             dfs2(edge[i].v,now);
39         }
40 }
41 int main()
42 {
43     ios::sync_with_stdio(false);
44     cin.tie(0);cout.tie(0);
45     cin>>n>>k;
46     memset(head,-1,sizeof(head));
47     for(int i=1;i<n;i++)
48     {
49         cin>>u1>>v1;
50         add(u1,v1),add(v1,u1);
51     }
52     for(int i=1;i<=n;i++)
53         cin>>c[i];
54     dfs1(1,0);dfs2(1,0);
55     for(int i=1;i<=n;i++)
56         for(int j=0;j<=k;j++)
57             ans[i]+=f[i][j];
58     for(int i=1;i<=n;i++)
59         cout<<ans[i]<<'\n';
60     return 0;
61 }