25. Sum of Distances in Tree

Hard : IMP

An undirected, connected tree with N nodes labelled 0...N-1 and N-1 edges are given.

The ith edge connects nodes edges[i][0] and edges[i][1] together.

Return a list ans, where ans[i] is the sum of the distances between node i 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: 
Here is a diagram of the given tree:
  0
 / \
1   2
   /|\
  3 4 5
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.

Solution: (Using Two DFS)

Approach

1. Solve it with node 0 as root.

2. Initial an array count, count[i] counts all nodes in the subtree i. Initial an array of res, res[i] counts sum of distance in subtree i.

3. 1st DFS call is similar to basic approach finding the distance and sum of each node from index 0. Update the count and res

   count[root] = sum(count[i]) + 1 res[root] = sum(res[i]) + sum(count[i])

4. 2nd DFS is the Dp approach When we move our root from parent to its child i, count[i] points get 1 closer to root, n - count[i] nodes get 1 further to root. res[i] = res[root] - count[i] + N - count[i]

class Solution
{
public:
    vector<int> count;
    vector<int> res;

    void dfs(vector<int> v[], vector<bool> &vis, int node)
    {
        if (!vis[node])
        {
            vis[node] = 1;

            for (int i = 0; i < v[node].size(); i++)
            {
                if (!vis[v[node][i]])
                {
                    dfs(v, vis, v[node][i]);
                    count[node] += count[v[node][i]];
                    res[node] += res[v[node][i]] + count[v[node][i]];
                }
            }
            count[node]++;
        }
        return;
    }

    void dfs2(vector<int> v[], vector<bool> &vis, int node)
    {
        if (!vis[node])
        {
            vis[node] = 1;

            for (int i = 0; i < v[node].size(); i++)
            {
                if (!vis[v[node][i]])
                {
                    res[v[node][i]] = res[node] - count[v[node][i]] + count.size() - count[v[node][i]];
                    dfs2(v, vis, v[node][i]);
                }
            }
        }
        return;
    }

    vector<int> sumOfDistancesInTree(int N, vector<vector<int>> &edges)
    {

        vector<int> v[N];
        vector<bool> vis(N, false);
        count.assign(N, 0);
        res.assign(N, 0);

        for (int i = 0; i < edges.size(); i++)
        {

            v[edges[i][0]].push_back(edges[i][1]);
            v[edges[i][1]].push_back(edges[i][0]);
        }

        dfs(v, vis, 0);

        vis.assign(N, 0);

        dfs2(v, vis, 0);

        return res;
    }
};