leetcode 834. Sum of distances in the tree

Title Description

Given an undirected 、 Connected tree . There is... In the tree n The marks are 0…n-1 And n-1 side .

Given integer n And an array edges , edges[i] = [ai, bi] Represents the node in the tree ai and bi There is a side between .

Return length is n Array of answer , among answer[i] It's the third in the tree i The sum of the distances between a node and all other nodes .
Example 1:

Input : n = 6, edges = [[0,1],[0,2],[2,3],[2,4],[2,5]]
Output : [8,12,6,10,10,10]
explain : The tree is as shown in the picture .
We can calculate dist(0,1) + dist(0,2) + dist(0,3) + dist(0,4) + dist(0,5)
That is to say 1 + 1 + 2 + 2 + 2 = 8. therefore ,answer[0] = 8, And so on .
Example 2:

Input : n = 1, edges = []
Output : [0]
Example 3:

Input : n = 2, edges = [[1,0]]
Output : [1,1]
Tips :
1 <= n <= 3 * 10⁴
edges.length == n - 1
edges[i].length == 2
0 <= ai, bi < n
ai != bi
The given input is guaranteed to be a valid tree

analysis

First consider using u A tree for roots ,u How to find the sum of distances to other nodes .j yes u The child node of ,dp[u] Express u To u Is the sum of the distances of other nodes in the subtree of the root . be dp[u] = sum(dp[j] + cnt[j]); among cnt[j] Said to j Is the total number of nodes in the subtree of the root .j Yes cnt[j] - 1 A descendant node ,u The distance to these nodes is only j It's a long way from them 1, That's more cnt[j] - 1, Plus u To j Distance of 1, Namely dp[j] + cnt[j] 了 . So we need dp[u], Only need to u Of all child nodes dp and cnt Value can be calculated .
The idea is ：

int dfs1(int u,int fa) {
        cnt[u] = 1;
        int res = 0;
        for(int i = h[u];~i;i = ne[i]) {
            int j = e[i];
            if(j == fa) continue;
            int m = dfs1(j,u);
            cnt[u] += m;
            res += m + dp[j];
        }
        dp[u] = res;
        return cnt[u];
    }

seek u To u The problem of the sum of the distances of other nodes in a subtree that is the root can at best be considered as medium difficulty , The reason why this question is difficult is hard Because we need to find the sum of the distances from all nodes in the tree to other nodes , If you call next for each node dfs, The time complexity is squared , The data scale of this question is 30000 , The scale of undirected graph is 60000 , The square level algorithm will obviously TLE.
Do it once dfs It has been worked out u Sum of distances to all nodes , In the u Is any node in the subtree of the root v,dp[v] Express v The sum of the distances to its descendant nodes , This is the information we have worked out . set up ans[u] Express u The sum of the distances to other nodes in the tree , obviously ans[u] = dp[u], But other nodes in the tree only find the sum of the distances to their descendants , You also need to add the sum of the distances to non descendant nodes .

Let's take this picture as an example , For the root node 0 Do Times dfs Find out 0 Sum of distances to other nodes , At the same time like dp[2] Express 2 To its descendant node 3 4 5 The sum of the distances . Obviously the ultimate ans[2] Is equal to dp[2] Plus 2 To 0 and 1 Distance of , set up 0 To except 2 The sum of the distances of the nodes outside this branch is s, that 2 To 2 Non descendant nodes of （0 and 1） The sum of the distances is s + 2, This is equivalent to 0 regard as 2 The child node of , And before dfs equally , Add the child node dp Sum of numbers cnt Count , But what I worked out before dp[0] Or call it ans[0] Statistics 0 Distance to all nodes and , So subtract 0 To 2 The sum of the distance of this branch is 0 To 1 The distance and , namely s = ans[0] - dp[2] - cnt[2], therefore ans[2] = dp[2] + s + n - cnt[2] = dp[2] + ans[0] - dp[2] -cnt[2] + n - cnt[2], among n - cnt[2] Express 2 The number of non descendant nodes of .
It seems that reasoning is troublesome , It's very simple ,2 The distance to all nodes is 2 The distance to the descendant node plus the distance to the non descendant node , The distance of non descendant nodes can pass through 2 The distance from the parent node of to all nodes minus to 2 The distance of all nodes on this branch , Add the number of non descendant nodes to find .
Summarize the idea of this topic , First, do the following for the root node dfs, Find the sum of the distances from each node to its descendant nodes and save it to dp In the array . Then do the root node again dfs, According to the ans Value and its own dp The value of each node is calculated from top to bottom ans value .

Code

class Solution {
public:
    static const int N = 60005;
    int idx = 0,e[N],h[N],ne[N];
    int tot = 0,cnt[N];
    vector<int> ans,dp;
    void add(int a,int b) {
        e[idx] = b,ne[idx] = h[a], h[a] = idx++;
    }
    int dfs1(int u,int fa) {
        cnt[u] = 1;
        int res = 0;
        for(int i = h[u];~i;i = ne[i]) {
            int j = e[i];
            if(j == fa) continue;
            int m = dfs1(j,u);
            cnt[u] += m;
            res += m + dp[j];
        }
        dp[u] = res;
        return cnt[u];
    }
    void dfs2(int u,int fa) {// Find out ans[u], Then recursively find the of its child nodes ans value 
        if(fa != -1)    ans[u] += ans[fa] - cnt[u] - dp[u] + tot - cnt[u];
        for(int i = h[u];~i;i = ne[i]) {
            int j = e[i];
            if(j != fa) {
                dfs2(j,u);
            }
        }
    }
    vector<int> sumOfDistancesInTree(int n, vector<vector<int>>& edges) {
        memset(h,-1,sizeof h);
        tot = n;
        int u;
        for(auto t : edges) add(t[0],t[1]),add(t[1],t[0]);// Construct an undirected graph , There is no need to find the root node 
        dp = vector<int>(n,0);
        dfs1(0,-1);
        ans = dp;// First add the distance to the descendant node 
        dfs2(0,-1);
        return ans;
    }
};