P3047 [USACO12FEB]附近的牛Nearby Cows

https://www.luogu.org/problemnew/show/P304 1 #include <bits/stdc++.h> 2 #define up(i,l,r) for(register int i = (l); i <= (r); ++i) 3 #define dn(i,l,r) for(register int i = (l); i >= (r); --i) 4 #define ll long long

 #define re register
 using namespace std;

 template <typename T> void in(T &x) {
     x = ; T f = ; char ch = getchar();
     while(!isdigit(ch)) {if(ch == '-') f = -; ch = getchar();}
     while( isdigit(ch)) {x =  * x + ch - ; ch = getchar();}
     x *= f;
 }

 template <typename T> void out(T x) {
     if(x < ) x = -x , putchar('-');
     if(x > ) out(x/);
     putchar(x% + );
 }

 int n,k;
 const int N = ;
 struct edge {
     int v,nxt;
 }e[N<<]; int tot,head[N];
 void add(int u,int v) {e[++tot] = (edge){v,head[u]}; head[u] = tot;}

 int f[][N][],c[N];

 void dfs1(int u,int fa) {
     up(i,,k)
     f[][u][i] = c[u];
     for(re int i = head[u]; i ; i = e[i].nxt) {
         int v = e[i].v;
         if(v == fa) continue;
         dfs1(v,u);
         up(i,,k)
         f[][u][i] += f[][v][i-];
     }
 }

 void dfs2(int u,int fa) {

     f[][u][] = c[u];
     up(i,,k) f[][u][i] = f[][fa][i-] - f[][u][i-] + f[][u][i];
     //别忘了+f[0][u][i]
     for(re int i = head[u]; i ; i = e[i].nxt) {
         int v = e[i].v;
         if(v == fa) continue;
         dfs2(v,u);
     }
 }

 int main() {
     freopen("input.txt","r",stdin);
     in(n); in(k);
     int u,v;
     up(i,,n-) {
         in(u); in(v);
         add(u,v); add(v,u);
     }
     up(i,,n) in(c[i]);
     dfs1(,);

     up(i,,k)
     f[][][i] = f[][][i-];
     //必须给0号节点赋初值
     dfs2(,);

     up(i,,n) out(f[][i][k]),putchar('\n');
     return ;
 }