"Algorithm Competition·Quick 300 Questions" One question per day: "Nearby Cow"

" Algorithm Competition: 300 Quick Questions " will be published in 2024 and is an auxiliary exercise book for "Algorithm Competition" .
All questions are placed in the self-built OJ New Online Judge .
Codes are given in three languages: C/C++, Java, and Python. The topics are mainly mid- to low-level topics and are suitable for entry-level and advanced students.

" Cows Nearby ", link: http://oj.ecustacm.cn/problem.php?id=1897

Question description

[Problem description] The farm consists of N points and N-1 two-way paths. Connect all points through N-1 paths.
There are C(i) cows at point i, and the cows sometimes move to different points through K paths.
The farmer wants to plant enough grass at each point i to feed the maximum number of cows M(i).
That is, for each point i, the farmer needs to consider the maximum number of cows that may reach the point i.
Request each M(i).
[Input format] Line 1: two integers, N and K, 1 <= N <= 100,000, 1 <= K <= 20.
Line 2 - Line N: Two integers u and v, indicating that there is a path between point u and point v, 1 <= u, v <= N.
Line N + 1 - Line 2N: Line N + i enters 1 integer to represent C(i), 0 <= C(i) <= 1000.
[Output format] Output N lines, and the i-th line outputs M(i).
【Input sample】

【Output sample】

answer

Simple summary of the question: A tree with n points and n-1 edges. Each point has a weight. For each node, find the sum m of the weights of all nodes that are no more than k away from it.
Consider the law of violence first. For any point i, directly traverse all points whose distance is no more than k, and find the sum of its weights mi. Encoding uses dfs, returns from each point dfs, depth is k. How much calculation is required? Assume that this tree is a full binary tree. If you take k steps from a point, you may reach $2^k$ points, when k=20, $2^{20}>100000$ , including all n points. Therefore, finding the m of one point alone is O(n), and finding the m of n points is $O(n^2)$ , timeout.
How to optimize? Calculating m separately for each point results in a large number of repeated calculations. For example, most of the two adjacent points u and v, the points within the distance k of the point u, and the points within the distance k of the point v are repeated, and only those different points need to be calculated.
Therefore, the idea and coding steps of this question are: (1) First calculate the sum of weights si of the subtree rooted at any point i, and do a DFS; (2) Then calculate the sum of weights starting from any point i and mi, which includes the sum of i's subtree weights si, and the weight of i's parent node direction. This calculation uses the previous results.
(1) Calculate the sum of weights of the subtrees of any point i.
Assume the state is dp[][], dp[i][j] represents the sum of the weights of the child nodes reaching j steps from i on the subtree with the i-th node as the root. Note that it is not the sum of the weights within j steps from i, but the sum of the weights of the child nodes that are j steps away. The code uses the dfs1() function to perform DFS once, and the dp[][] of each point can be calculated.
Note that dp[][] has a new meaning in the following calculations: dp[i][j] represents the sum of the weights of nodes arriving in j steps starting from i. It includes not only the j-th step node on the subtree of i, but also the j-th step node in the direction of the parent node.
(2) Calculate the sum m of the weights of all nodes within the distance k from point v, as shown in the figure below.
Insert image description here

It includes two parts:
1) The node on the subtree where v goes downwards. The weight of this part is equal to the dp[v][j] calculated previously in dfs1(), 0≤j≤k.
2) How to calculate the sum of the weights of the nodes within the distance k that v moves upward? One step up from v is the parent node u. The dp[u][j-1] corresponding to u is the sum of the weights of the nodes at the k-1th step away from u. It includes the dotted line (1) in the figure, The sum of the weights of several nodes at the end points of the arrows in (2) and (3). When calculating the sum of the weights of v in the u direction, you should add (2) and (3) instead of (1), that is, remove dp[v][j- from dp[u][j-1] 2], that is, dp[u][j-1] - dp[v][j-2].
In the code dfs2() function, use tot[] to record the answer, tot[u] is the sum of weights within the distance k from the u-th point. dp[u][i] in dfs2() has been updated to a new meaning. In line 20, dp[u][i] is accumulated, 0≤j≤k, and tot[u] can be calculated.
Lines 21-26 calculate and update dp[][] to be the sum of weights under the new meaning. Be careful not to forget to update dp[v][1] and add the weight dp[u][0] of v's parent node u on line 24.
dfs2() has two keys.
1) In line 23, j goes from k to 2 and loops in reverse. dp[v][j] += dp[u][j - 1] - dp[v][j - 2], the dp[][] on the left is the new meaning, and the dp[][] on the right is dfs1( ) is the old meaning when calculating, j looping backwards can avoid destroying this relationship.
2) Line 25 continues dfs2() at the end, that is, it continues DFS after updating dp[][] previously. When entering dfs2() for the first time, u is the root node 1 of the entire tree. It has no parent node. Calculate tot[1] and directly accumulate dp[u][i]. When continuing to calculate the tot[] of the child node later, you need to calculate the weights of the two parts according to the instructions in the figure above.
The computational complexity of both dfs1() and dfs1() is O(n).

[Key Points] Tree DP.

C++ code

#include<bits/stdc++.h>
using namespace std;
const int N = 100010;
const int K = 22;
int n, k;
int c[N];
vector<int> e[N];                   //存图
int dp[N][K];
void dfs1(int u, int fa) {
    
              //dp[i][j]:i往子树走j步，第j步子节点权值之和
    dp[u][0] = c[u];                //先加上自己的权值
    for(auto v : e[u]){
    
    
        if(v == fa) continue;
        dfs1(v, u);                 //注意：先dfs，计算出子节点的dp[][]，回溯后带回
        for(int j=1; j<=k; j++)     //从第j=1步开始，一步一步往下走并计算
            dp[u][j] += dp[v][j - 1];
    }
}
int tot[N];
void dfs2(int u, int fa) {
    
    
    for(int i=0; i<=k; i++)  tot[u] += dp[u][i]; //第 1)部分：先累加u子树上，k步内的权值之和        
    for(auto v : e[u]){
    
                              //第 2)部分：然后计算u的子节点的权值之和
        if(v == fa) continue;
        for(int j=k; j>=2; j--)  dp[v][j] += dp[u][j - 1] - dp[v][j - 2];
        dp[v][1] += dp[u][0];                    //v往父节点u走一步，加上u的权值
        dfs2(v, u);
    }
}
int main(){
    
    
    cin>>n>>k;
    for(int i=1; i<n; i++) {
    
    
        int u, v; scanf("%d%d", &u, &v);
        e[u].push_back(v);
        e[v].push_back(u);
    }
    for(int i=1; i<=n; i++)  scanf("%d", &c[i]);
    dfs1(1, 1);     //以1为根，求每个点往子节点方向走k步的m值
    dfs2(1, 1);     //仍然从根1出发，求每个点的权值之和
    for(int i=1; i<=n; i++)  printf("%d\n", tot[i]);
    return 0;
}