834

834. Sum of Distances in Tree (Hard)

There is an undirected connected tree with n nodes labeled from 0 to n - 1 and n - 1 edges.

You are given the integer n and the array edges where edges[i] = [ai, bi] indicates that there is an edge between nodes ai and bi in the tree.

Return an array answer of length n where answer[i] is the sum of the distances between the ith node in the tree 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]
Explanation: The tree is shown above.
We can see that dist(0,1) + dist(0,2) + dist(0,3) + dist(0,4) + dist(0,5)
equals 1 + 1 + 2 + 2 + 2 = 8.
Hence, 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]

 

Constraints:

• 1 <= n <= 3 * 104
• edges.length == n - 1
• edges[i].length == 2
• 0 <= ai, bi < n
• ai != bi
• The given input represents a valid tree.

Companies:
Google

Related Topics:
Dynamic Programming, Tree, Depth-First Search, Graph

Similar Questions:

• Distribute Coins in Binary Tree (Medium)

Solution 1. DFS

// OJ: https://leetcode.com/problems/sum-of-distances-in-tree/
// Author: github.com/lzl124631x
// Time: O(N)
// Space: O(N)
class Solution {
    vector<vector<int>> G;
    vector<int> cnt, ans;
    int root = 0;
    int getSum(int u, int prev = -1, int d = 0) {
        int ans = d;
        for (int v : G[u]) {
            if (v == prev) continue;
            ans += getSum(v, u, d + 1);
        }
        return ans;
    }
    int updateCount(int u, int prev = -1) {
        int ans = 0;
        for (int v : G[u]) {
            if (v == prev) continue;
            ans += updateCount(v, u);
        }
        return cnt[u] = ans + 1;
    }
    void updateAns(int u, int prev = -1) {
        if (prev != -1) ans[u] = ans[prev] + cnt[root] - 2 * cnt[u];
        for (int v : G[u]) {
            if (v == prev) continue;
            updateAns(v, u);
        }
    }
public:
    vector<int> sumOfDistancesInTree(int n, vector<vector<int>>& E) {
        if (n == 1) return {0};
        G.assign(n, vector<int>());
        cnt.assign(n, 0);
        ans.assign(n, 0);
        vector<int> indegree(n);
        for (auto &e : E) {
            int u = e[0], v = e[1];
            G[u].push_back(v);
            G[v].push_back(u);
            ++indegree[u];
            ++indegree[v];
        }
        for (; indegree[root] != 1; ++root); // pick a leaf node as root
        ans[root] = getSum(root); // get the sum of distances of this node `root`
        updateCount(root); // update the cnt[i] which is the number of nodes in the subtree with `i` as the subtree root
        updateAns(root); // update the answer: ans[u] = ans[prev] + cnt[root] - 2 * cnt[u]
        return ans;
    }
};

TODO

Simplify code according to https://leetcode.com/problems/sum-of-distances-in-tree/solution/