Sum of Distances in Tree


Problem

There is an undirected connected tree with nn nodes labeled from 00 to n1n - 1 and n1n - 1 edges.

You are given the integer nn and the array edges where edgesi=[ai,bi]edges_i = [a_i, b_i] indicates that there is an edge between nodes aia_i and bib_i in the tree.

Return an array answeranswer of length nn where answerianswer_i is the sum of the distances between the ithi^{th} node in the tree and all other nodes.

Example

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

  • 1n3×1041 \leq n \leq 3 \times 10^4
  • edges.length=n1edges.length = n - 1
  • edges[i].length=2edges[i].length = 2
  • 0ai,bi<n0 \leq a_i, b_i < n
  • aibia_i \neq b_i

Submit your solution at here

Solution

Intuition

The first naive solution came to my mind is DFS each node and count. This brute force will result in O(n2)O(n^2) and obviously TLE.

But what if we have already have the answer for node ii, can we calculate answer for all children of ii fast?

Approach

Actually we can do it, let’s DFS twice:

  • First pass, for each node ii calculate sum of all path to all children of ii
    • After first pass, the value on the root node already is the answer
  • Second pass, for each node ii with parent pp that has already has answer, calculate sum of all path to the other node that on the other side (is not children) of ii

Complexity

  • Time complexity: O(n)O(n)
  • Space complexity: O(n)O(n)

Code

The implementation is heavy, looks daunting but actually not too complex 😀

class Solution {
public:
    vector<int> parent;
    vector<int> childCount;
    vector<vector<int>> children;
    vector<vector<int>> adj;
    void buildTree(vector<vector<int>>& adj) {
        vector<bool> vis(parent.size(), false);
        queue<int> q;
        q.push(0);
        vis[0] = true;
        while(!q.empty()) {
            auto u = q.front();
            q.pop();
            for(auto v: adj[u]) {
                if(vis[v]) {
                    continue;
                }
                vis[v] = true;
                q.push(v);
                parent[v] = u;
                children[u].push_back(v);
            }
        }
    }
    void buildDistBottomUp(int root, vector<int>& score) {
        int p = parent[root];
        if(p != -1) {
            int n = parent.size();
            int otherSideScore = score[p] - score[root] - childCount[root] - 1;
            int otherCount = n - childCount[root] - 1;
            //cout<<"build "<<root<<" current score = "<<dist[root]<<", other side score = "<<otherSideScore<<" other count = "<<otherCount<<endl;
            score[root] += otherSideScore + otherCount;
        }
        for(auto child: children[root]) {
            buildDistBottomUp(child, score);
        }
    }
    void buildDistTopDown(int root, vector<int>& score) {
        for(auto child: children[root]) {
            buildDistTopDown(child, score);
            score[root] += score[child]+childCount[child]+1;
        }
    }
    void countChildren(int root) {
        for(auto child: children[root]) {
            countChildren(child);
            childCount[root] += childCount[child]+1;
        }
    }
    vector<int> sumOfDistancesInTree(int n, vector<vector<int>>& edges) {
        parent.resize(n,-1);
        childCount.resize(n, 0);
        children.resize(n);
        adj.resize(n);
        for(auto e: edges) {
            adj[e[0]].push_back(e[1]);
            adj[e[1]].push_back(e[0]);
        }
        vector<int> ret(n, 0);
        buildTree(adj);
        countChildren(0);
        buildDistTopDown(0, ret);
        buildDistBottomUp(0, ret);
        return ret;
    }
};
Made with love using Svelte, Three.js