2023-07-16 One question per day
1. Topic number
834. 树中距离之和
2. Topic link
Click to jump to the topic location
3. Topic description
Given an undirected, connected tree. There are n nodes labeled 0…n-1 and n-1 edges in the tree.
Given an integer n and an array edges, edges[i] = [ai, bi] means that there is an edge between nodes ai and bi in the tree.
Returns an array answer of length n, where answer[i] is the sum of the distances between the i-th node and all other nodes in the tree.
hint:
- 1 <= n <= 3 * 104
- edges.length == n - 1
- edges[i].length == 2
- 0 <= ai, bi < n
- ai != bi
- The given input is guaranteed to be a valid tree
4. Problem solving code
class Solution {
int dp[30004];//表示dp[u] 表示以u为根的子树,它的所有子节点到它的距离之和
int sz[30004];//表示以u为根的子树的节点数量
vector<int> ans;//表示最终的答案
vector<int>Edge[30004];
void dfs1(int u, int f) {
sz[u] = 1;
dp[u] = 0;
for (int i = 0; i < Edge[u].size(); ++i) {
int v = Edge[u][i];
if (v == f) {
continue;
}
dfs1(v, u);
dp[u] += dp[v] + sz[v];
sz[u] += sz[v];
}
}
void dfs2(int u, int f) {
ans[u] = dp[u];
for (int i = 0; i < Edge[u].size(); ++i) {
int v = Edge[u][i];
if (v == f) {
continue;
}
int pu = dp[u], pv = dp[v];
int su = sz[u], sv = sz[v];
dp[u] -= dp[v] + sz[v];
sz[u] -= sz[v];
dp[v] += dp[u] + sz[u];
sz[v] += sz[u];
dfs2(v, u);
dp[u] = pu, dp[v] = pv;
sz[u] = su, sz[v] = sv;
}
}
public:
vector<int> sumOfDistancesInTree(int n, vector<vector<int>>& edges) {
for(int i = 0; i < n; ++i){
Edge[i].clear();
}
for(int i = 0; i < n - 1; ++i){
int x = edges[i][0];
int y = edges[i][1];
Edge[x].push_back(y);
Edge[y].push_back(x);
}
memset(dp, 0, sizeof(dp));
memset(sz, 0, sizeof(sz));
ans.resize(n);
dfs1(0, -1);
dfs2(0, -1);
return ans;
}
};
Five, problem-solving ideas
(1) Use tree dp to solve this problem.
(2) Set the initial state first, dp[u] represents the subtree rooted at u, and the sum of the distances from all its child nodes to it. sz[u] represents the number of nodes in the subtree rooted at u. vector ans represents the final answer.
(3) Then we need to ask a question, what we need is how to get dp[u]. Node u has x 1 , x 2 , x 3 child nodes, then dp[u] should be equal to 3+dp[x 1 ]+dp[x 2 ]+dp[x 3 ]. Then sz[u] means sz[x 1 ] + sz[x 2 ] + sz[x 3 ], this is the work done in dfs1. Of course all sz should be 1 at the beginning. We follow this way of thinking n times, then the problem will definitely be solved, but the time complexity is very high, can it be optimized?
(4) After a tree dynamic programming, we actually obtained the answer dp of the subtree rooted at each node in the tree rooted at u, and we can use these existing information to optimize the time complexity. Originally, what we were looking for was to use u as the root node and v as the child node. We will replace it later, so v is the root node and u is the child node. Except for the values of dp[u] and dp[v], other values do not need to be changed. First u subtracts v's contribution. v is to be added to u's contribution. In this way, we have realized the operation of changing the root.