Open In App

Sum of length of paths from every node to all other nodes using Tree Rerooting technique

Improve
Improve
Like Article
Like
Save
Share
Report

For every node in an undirected tree, find the sum of length of paths from it to all other nodes, using Tree Rerooting technique.

Rerooting is a dynamic programming technique in trees. Using this a property dp[i] on tree is calculated, when the tree is rooted at i.

Example: 

Input: Consider the tree shown in the image:

tree rooted at 1 (arrow represent paths from root to other nodes)

Output: 12 9 17 14 10 15 15

Explanation: for node i : i1 + i2 + i3 . . . in, where ij is length of path from i to j
For node 1 : 0 + 1 + 1 + 2 + 2 + 3 + 3 = 12
For node 2 : 1 + 0 + 2 + 1 + 1 + 2 + 2 = 9
For node 3 : 1 + 2 + 0 + 3 + 3 + 4 + 4 = 17
For node 4 : 2 + 1 + 3 + 0 + 2 + 3 + 3 = 14
For node 5 : 2 + 1 + 3 + 2 + 0 + 1 + 1 = 10
For node 6 : 3 + 2 + 4 + 3 + 1 + 0 + 2 = 15
For node 7 : 3 + 2 + 4 + 3 + 1 + 2 + 0 = 15

 

Naive approach: This approach is based on the following observation of dynamic programming.

Consider at following figure to understand the transition in dynamic programming:

Illustration:

tree dp for sum of path lengths

Note: dp[node] in figure denotes the sum of paths from node to all its subtrees.

Now Consider child ‘b’ of node ‘a’

  1. Size of subtree ‘b’, size[b] and the sum of lengths of paths for subtree ‘b’ is given as dp[b].
  2. Consider all the paths which start at ‘a’ and end at some node in subtree of ‘b’.
  3. Every path can be broken down as follows: path(a, x) = edge(a, b) + path(b, x)
  4. As lengths of all paths of form path(b, x) is already covered in dp[b], only the edge(a, b) needs to be added for all nodes ‘x’ in subtree of ‘b’.
  5. That means size[b] needs to be added to dp[b], hence the contribution of subtree ‘b’ to dp[a] is (dp[b] + size[b]).

The transitions are as following:

Transition :
for (child of node):
    dp[node] += dp[child] + size[child]

Follow the steps mentioned below to implement the approach:

  1. For each node do the following :
    • Root the tree at this node, and find the dp as described above
  2. As the tree is rooted at ‘node’, all the other nodes will be in its subtree. So dp[node] will be the required answer for ‘node’.

C++




// C++ code to implement above approach
#include <bits/stdc++.h>
using namespace std;
 
// Dfs computes dp, answer for each node
// with respect to its subtree it also
// computes size of each subtree
void dfs(int node, int par,
         vector<vector<int> >& g,
         vector<int>& size, vector<int>& dp)
{
    // Initialise given subtree with dp = 0
    // [as there is no paths currently] and
    // size 1, because there is only
    // one node in subtree
    size[node] = 1;
    dp[node] = 0;
    for (auto nebr : g[node]) {
         
        // For every neighbour of node
        // which is not its parent
        // 1. compute size and dp for
        // nebr by dfs
        // 2. update size and dp for node,
        // based on nebr
        // See explanation to understand
        // the dp transition
        if (nebr != par) {
            dfs(nebr, node, g, size, dp);
            size[node] += size[nebr];
            dp[node] += dp[nebr] +
                size[nebr];
        }
    }
}
 
// Creates a edge between a and b,
// given graph g
void edge(int a, int b,
          vector<vector<int> >& g)
{
    // Convert into 0-based indexing
    a--;
    b--;
 
    // Push b to adjacency list of a
    // and vice versa because given
    // tree is undirected
    g[a].push_back(b);
    g[b].push_back(a);
}
 
// Function to get the sum of paths
vector<int> pathSum(vector<vector<int> > &g,
                   int N)
{
    vector<int> dp(N), ans(N), size(N);
    // For root 'r'
    // 1. compute dp for tree rooted at 'r'
    // 2. as all nodes belong to some
    // subtree of root, answer will be
    // equal to dp
    for (int r = 0; r < N; ++r) {
        dfs(r, -1, g, size, dp);
        ans[r] = dp[r];
    }
    return ans;
}
 
// Driver code
int main()
{
    int N = 7;
 
    vector<vector<int> > g(N);
 
    edge(1, 2, g);
    edge(1, 3, g);
    edge(2, 4, g);
    edge(2, 5, g);
    edge(5, 6, g);
    edge(5, 7, g);
    vector<int> res = pathSum(g, N);
     
    for (int i = 0; i < N; ++i) {
        cout << res[i] << " ";
    }
    cout << endl;
    return 0;
}


Java




// Java code to implement above approach
import java.util.*;
 
class GFG
{
   
  static int N = 7;
  static int dp[] = new int[N];
  static int ans[] = new int[N];
  static int size[] = new int[N];
 
  // Dfs computes dp, answer for each node
  // with respect to its subtree it also
  // computes size of each subtree
  static void dfs(int node, int par, Vector<Integer> []g)
  {
 
    // Initialise given subtree with dp = 0
    // [as there is no paths currently] and
    // size 1, because there is only
    // one node in subtree
    size[node] = 1;
    dp[node] = 0;
    for (int nebr : g[node]) {
 
      // For every neighbour of node
      // which is not its parent
      // 1. compute size and dp for
      // nebr by dfs
      // 2. update size and dp for node,
      // based on nebr
      // See explanation to understand
      // the dp transition
      if (nebr != par) {
        dfs(nebr, node, g);
        size[node] += size[nebr];
        dp[node] += dp[nebr] +
          size[nebr];
      }
    }
  }
 
  // Creates a edge between a and b,
  // given graph g
  static void edge(int a, int b,
                   Vector<Integer> [] g)
  {
    // Convert into 0-based indexing
    a--;
    b--;
 
    // Push b to adjacency list of a
    // and vice versa because given
    // tree is undirected
    g[a].add(b);
    g[b].add(a);
  }
 
  // Function to get the sum of paths
  static int[] pathSum(Vector<Integer> []g)
  {
 
    // For root 'r'
    // 1. compute dp for tree rooted at 'r'
    // 2. as all nodes belong to some
    // subtree of root, answer will be
    // equal to dp
    for (int r = 0; r < N; ++r) {
      dfs(r, -1, g);
      ans[r] = dp[r];
    }
    return ans;
  }
 
  // Driver code
  public static void main(String[] args)
  {
 
    Vector<Integer> []g = new Vector[N];
    for (int i = 0; i < g.length; i++)
      g[i] = new Vector<Integer>();
    edge(1, 2, g);
    edge(1, 3, g);
    edge(2, 4, g);
    edge(2, 5, g);
    edge(5, 6, g);
    edge(5, 7, g);
    int[] res = pathSum(g);
 
    for (int i = 0; i < N; ++i) {
      System.out.print(res[i]+ " ");
    }
    System.out.println();
  }
}
 
// This code is contributed by 29AjayKumar


C#




// C# code to implement above approach
using System;
using System.Collections.Generic;
class GFG
{
 
    static int N = 7;
    static int[] dp = new int[N];
    static int[] ans = new int[N];
    static int[] size = new int[N];
 
    // Dfs computes dp, answer for each node
    // with respect to its subtree it also
    // computes size of each subtree
    static void dfs(int node, int par, List<int>[] g)
    {
 
        // Initialise given subtree with dp = 0
        // [as there is no paths currently] and
        // size 1, because there is only
        // one node in subtree
        size[node] = 1;
        dp[node] = 0;
        foreach (int nebr in g[node])
        {
 
            // For every neighbour of node
            // which is not its parent
            // 1. compute size and dp for
            // nebr by dfs
            // 2. update size and dp for node,
            // based on nebr
            // See explanation to understand
            // the dp transition
            if (nebr != par)
            {
                dfs(nebr, node, g);
                size[node] += size[nebr];
                dp[node] += dp[nebr] +
                  size[nebr];
            }
        }
    }
 
    // Creates a edge between a and b,
    // given graph g
    static void edge(int a, int b,
                     List<int>[] g)
    {
        // Convert into 0-based indexing
        a--;
        b--;
 
        // Push b to adjacency list of a
        // and vice versa because given
        // tree is undirected
        g[a].Add(b);
        g[b].Add(a);
    }
 
    // Function to get the sum of paths
    static int[] pathSum(List<int>[] g)
    {
 
        // For root 'r'
        // 1. compute dp for tree rooted at 'r'
        // 2. as all nodes belong to some
        // subtree of root, answer will be
        // equal to dp
        for (int r = 0; r < N; ++r)
        {
            dfs(r, -1, g);
            ans[r] = dp[r];
        }
        return ans;
    }
 
    // Driver code
    public static void Main()
    {
 
        List<int>[] g = new List<int>[N];
        for (int i = 0; i < g.Length; i++)
            g[i] = new List<int>();
        edge(1, 2, g);
        edge(1, 3, g);
        edge(2, 4, g);
        edge(2, 5, g);
        edge(5, 6, g);
        edge(5, 7, g);
        int[] res = pathSum(g);
 
        for (int i = 0; i < N; ++i)
        {
            Console.Write(res[i] + " ");
        }
        Console.WriteLine();
    }
}
 
// This code is contributed by Saurabh Jaiswal


Python3




## Python program for the above approach:
 
## Dfs computes dp, answer for each node
## with respect to its subtree it also
## computes size of each subtree
def dfs(node, par, g, size, dp):
    ## Initialise given subtree with dp = 0
    ## [as there is no paths currently] and
    ## size 1, because there is only
    ## one node in subtree
    size[node] = 1
    dp[node] = 0
    for nebr in g[node]:
         
        ## For every neighbour of node
        ## which is not its parent
        ## 1. compute size and dp for
        ## nebr by dfs
        ## 2. update size and dp for node,
        ## based on nebr
        ## See explanation to understand
        ## the dp transition
        if (nebr != par):
            dfs(nebr, node, g, size, dp)
            size[node] += size[nebr]
            dp[node] += (dp[nebr] + size[nebr])
 
## Creates a edge between a and b,
## given graph g
def edge(a, b, g):
    ## Convert into 0-based indexing
    a-=1
    b-=1
 
    ## Push b to adjacency list of a
    ## and vice versa because given
    ## tree is undirected
    g[a].append(b);
    g[b].append(a);
 
## Function to get the sum of paths
def pathSum(g, N):
    dp = [0]*N
    ans = [0]*N
    size = [0]*N
    ## For root 'r'
    ## 1. compute dp for tree rooted at 'r'
    ## 2. as all nodes belong to some
    ## subtree of root, answer will be
    ## equal to dp
    for r in range(0, N):
        dfs(r, -1, g, size, dp)
        ans[r] = dp[r]
    return ans
 
 
 
## Driver code
if __name__=='__main__':
 
    N = 7
     
    g = [0]*N
    for i in range(0, N):
        g[i] = []
     
    edge(1, 2, g)
    edge(1, 3, g)
    edge(2, 4, g)
    edge(2, 5, g)
    edge(5, 6, g)
    edge(5, 7, g)
    res = pathSum(g, N)
     
    for i in range(0, N):
        print(res[i], end=' ')
    print("")


Javascript




// Javascript code for the above approach
  
 // Dfs computes dp, answer for each node
  // with respect to its subtree it also
  // computes size of each subtree
function dfs(node, par, g, size, dp)
{
 
    // Initialise given subtree with dp = 0
    // [as there is no paths currently] and
    // size 1, because there is only
    // one node in subtree
    size[node] = 1;
    dp[node] = 0;
    for (let nebr of g[node])
    {
     
        // For every neighbour of node
        // which is not its parent
        // 1. compute size and dp for
        // nebr by dfs
        // 2. update size and dp for node,
        // based on nebr
        // See explanation to understand
        // the dp transition
        if (nebr !== par) {
            dfs(nebr, node, g, size, dp);
            size[node] += size[nebr];
            dp[node] += dp[nebr] + size[nebr];
        }
    }
}
 
function edge(a, b, g) {
    // Convert into 0-based indexing
    a -= 1;
    b -= 1;
 
    // Push b to adjacency list of a
    // and vice versa because given
    // tree is undirected
    g[a].push(b);
    g[b].push(a);
}
 
function pathSum(g, N) {
    let dp = new Array(N).fill(0);
    let ans = new Array(N).fill(0);
    let size = new Array(N).fill(0);
     
    // For root 'r'
    // 1. compute dp for tree rooted at 'r'
    // 2. as all nodes belong to some
    // subtree of root, answer will be
    // equal to dp
    for (let r = 0; r < N; r++) {
        dfs(r, -1, g, size, dp);
        ans[r] = dp[r];
    }
    return ans;
}
 
    let N = 7;
    let g = new Array(N).fill(null).map(() => []);
    edge(1, 2, g);
    edge(1, 3, g);
    edge(2, 4, g);
    edge(2, 5, g);
    edge(5, 6, g);
    edge(5, 7, g);
    let res = pathSum(g, N);
 
    for (let i = 0; i < N; i++) {
        console.log(res[i]+" ");
    }
   
// This code is contributed by ik_9


Output

12 9 17 14 10 15 15 

Time Complexity: O(N2), where N is the number of nodes.
Auxiliary Space: O(N)

Rerooting Approach: The solution can be further optimised  by calculating answer for one root and rerooting the tree everytime to calculate for other nodes.

Look at following figure to understand the rerooting concept:

 

rerooting tree from 1 to 2

 

Note : in above figure, the edges are undirected, the arrows only represent the paths from root to other nodes

  • In given figure there is rerooting from ‘1’ to ‘2’ using the following approach
    • remove 2 from children of 1
    • add 1 to children of 2

As there is no recomputation for every node and rerooting takes only O(1), the overall time complexity also decreases.

C++




// C++ code to implement above approach
#include <bits/stdc++.h>
using namespace std;
 
// The function dfs0 computes dp,
// answer for each node with respect to
// its subtree it also computes
// size of each subtree
void dfs0(int node, int par,
          vector<vector<int> >& g,
          vector<int>& dp, vector<int>& size)
{
    // Initialise given subtree with dp = 0
    // as there is no paths currently and
    // size 1, because there is only
    // one node in subtree
    dp[node] = 0;
    size[node] = 1;
    for (auto nebr : g[node]) {
         
        // For every neighbour of node
        // which is not its parent
        // 1. compute size and dp for
        // nebr by dfs
        // 2. update size and dp for node,
        // based on nebr
        // See explanation to understand
        // the dp transition
        if (par != nebr) {
            dfs0(nebr, node, g, dp, size);
            size[node] += size[nebr];
            dp[node] += size[nebr] +
                dp[nebr];
        }
    }
}
 
// Rerooting the tree from 'from' to 'to'
void reroot(int from, int to,
            vector<int>& dp,
            vector<int>& size)
{
    // 'to' is no longer a child of 'from'
    dp[from] -= size[to] + dp[to];
    size[from] -= size[to];
     
    // 'from' is now a child of 'to'
    size[to] += size[from];
    dp[to] += size[from] + dp[from];
}
 
void dfs1(int node, int par,
          vector<vector<int> >& g,
          vector<int>& dp, vector<int>& ans,
          vector<int>& size)
{
    // Current dfs considers 'node' as root
    // so currently dp[node]
    // will be the answer
    ans[node] = dp[node];
 
    // For all neighbours which are
    // not parent of node
    for (auto nebr : g[node]) {
        if (par != nebr) {
            // Reroot the tree to 'nebr'
            reroot(node, nebr, dp, size);
 
            // Compute ans for 'nebr'
            // as a root of tree with dfs
            dfs1(nebr, node, g, dp, ans,
                 size);
 
            // reroot the tree back
            // to 'node'
            reroot(nebr, node, dp, size);
        }
    }
}
 
// Creates a edge between a and b,
// given graph g
void edge(int a, int b,
          vector<vector<int> >& g)
{
    // Convert into 0-based indexing
    a--;
    b--;
 
    // push b to adjacency list
    // of a and vice versa
    // because given tree is undirected
    g[a].push_back(b);
    g[b].push_back(a);
}
 
// Function to calculate sum of paths
vector<int> pathSum(vector<vector<int> > &g,
                   int N)
{
    vector<int> dp(N), ans(N), size(N);
     
    // Compute answer for each subtree
    // with tree rooted at 0
    dfs0(0, -1, g, dp, size);
     
    // Compute answer for each node
    // as root of tree, rerooting
    dfs1(0, -1, g, dp, ans, size);
    return ans;
}
 
// Driver code
int main()
{
    int N = 7;
    vector<vector<int> > g(N);
 
    edge(1, 2, g);
    edge(1, 3, g);
    edge(2, 4, g);
    edge(2, 5, g);
    edge(5, 6, g);
    edge(5, 7, g);
     
    vector<int> res = pathSum(g, N);
    for (int i = 0; i < N; ++i) {
        cout << res[i] << " ";
    }
    cout << endl;
    return 0;
}


Java




// Java code to implement above approach
import java.util.*;
 
class GFG{
 
  static int N = 7;
  static int dp[] = new int[N];
  static int ans[] = new int[N];
  static int size[] = new int[N];
 
  // The function dfs0 computes dp,
  // answer for each node with respect to
  // its subtree it also computes
  // size of each subtree
  static void dfs0(int node, int par, Vector<Integer> []g)
  {
     
    // Initialise given subtree with dp = 0
    // as there is no paths currently and
    // size 1, because there is only
    // one node in subtree
    dp[node] = 0;
    size[node] = 1;
    for (int nebr : g[node]) {
 
      // For every neighbour of node
      // which is not its parent
      // 1. compute size and dp for
      // nebr by dfs
      // 2. update size and dp for node,
      // based on nebr
      // See explanation to understand
      // the dp transition
      if (par != nebr) {
        dfs0(nebr, node, g);
        size[node] += size[nebr];
        dp[node] += size[nebr] +
          dp[nebr];
      }
    }
  }
 
  // Rerooting the tree from 'from' to 'to'
  static void reroot(int from, int to)
  {
    // 'to' is no longer a child of 'from'
    dp[from] -= size[to] + dp[to];
    size[from] -= size[to];
 
    // 'from' is now a child of 'to'
    size[to] += size[from];
    dp[to] += size[from] + dp[from];
  }
 
  static void dfs1(int node, int par,Vector<Integer> []g)
  {
    // Current dfs considers 'node' as root
    // so currently dp[node]
    // will be the answer
    ans[node] = dp[node];
 
    // For all neighbours which are
    // not parent of node
    for (int nebr : g[node]) {
      if (par != nebr) {
        // Reroot the tree to 'nebr'
        reroot(node, nebr);
 
        // Compute ans for 'nebr'
        // as a root of tree with dfs
        dfs1(nebr, node, g);
 
        // reroot the tree back
        // to 'node'
        reroot(nebr, node);
      }
    }
  }
 
  // Creates a edge between a and b,
  // given graph g
  static void edge(int a, int b,
                   Vector<Integer> [] g)
  {
    // Convert into 0-based indexing
    a--;
    b--;
 
    // push b to adjacency list
    // of a and vice versa
    // because given tree is undirected
    g[a].add(b);
    g[b].add(a);
  }
 
  // Function to calculate sum of paths
  static int[] pathSum(Vector<Integer> []g)
  {
 
 
    // Compute answer for each subtree
    // with tree rooted at 0
    dfs0(0, -1, g);
 
    // Compute answer for each node
    // as root of tree, rerooting
    dfs1(0, -1, g);
    return ans;
  }
 
  // Driver code
  public static void main(String[] args)
  {
    int N = 7;
    Vector<Integer> []g = new Vector[N];
    for (int i = 0; i < g.length; i++)
      g[i] = new Vector<Integer>();
 
    edge(1, 2, g);
    edge(1, 3, g);
    edge(2, 4, g);
    edge(2, 5, g);
    edge(5, 6, g);
    edge(5, 7, g);
    int[] res = pathSum(g);
 
    for (int i = 0; i < N; ++i) {
      System.out.print(res[i]+ " ");
    }
    System.out.println();
  }
}
 
// This code contributed by shikhasingrajput


C#




// C# code to implement above approach
using System;
using System.Collections.Generic;
 
public class GFG{
 
  static int N = 7;
  static int []dp = new int[N];
  static int []ans = new int[N];
  static int []size = new int[N];
 
  // The function dfs0 computes dp,
  // answer for each node with respect to
  // its subtree it also computes
  // size of each subtree
  static void dfs0(int node, int par, List<int> []g)
  {
 
    // Initialise given subtree with dp = 0
    // as there is no paths currently and
    // size 1, because there is only
    // one node in subtree
    dp[node] = 0;
    size[node] = 1;
    foreach (int nebr in g[node]) {
 
      // For every neighbour of node
      // which is not its parent
      // 1. compute size and dp for
      // nebr by dfs
      // 2. update size and dp for node,
      // based on nebr
      // See explanation to understand
      // the dp transition
      if (par != nebr) {
        dfs0(nebr, node, g);
        size[node] += size[nebr];
        dp[node] += size[nebr] +
          dp[nebr];
      }
    }
  }
 
  // Rerooting the tree from 'from' to 'to'
  static void reroot(int from, int to)
  {
    // 'to' is no longer a child of 'from'
    dp[from] -= size[to] + dp[to];
    size[from] -= size[to];
 
    // 'from' is now a child of 'to'
    size[to] += size[from];
    dp[to] += size[from] + dp[from];
  }
 
  static void dfs1(int node, int par,List<int> []g)
  {
    // Current dfs considers 'node' as root
    // so currently dp[node]
    // will be the answer
    ans[node] = dp[node];
 
    // For all neighbours which are
    // not parent of node
    foreach (int nebr in g[node]) {
      if (par != nebr) {
        // Reroot the tree to 'nebr'
        reroot(node, nebr);
 
        // Compute ans for 'nebr'
        // as a root of tree with dfs
        dfs1(nebr, node, g);
 
        // reroot the tree back
        // to 'node'
        reroot(nebr, node);
      }
    }
  }
 
  // Creates a edge between a and b,
  // given graph g
  static void edge(int a, int b,
                   List<int> [] g)
  {
    // Convert into 0-based indexing
    a--;
    b--;
 
    // push b to adjacency list
    // of a and vice versa
    // because given tree is undirected
    g[a].Add(b);
    g[b].Add(a);
  }
 
  // Function to calculate sum of paths
  static int[] pathSum(List<int> []g)
  {
 
 
    // Compute answer for each subtree
    // with tree rooted at 0
    dfs0(0, -1, g);
 
    // Compute answer for each node
    // as root of tree, rerooting
    dfs1(0, -1, g);
    return ans;
  }
 
  // Driver code
  public static void Main(String[] args)
  {
    int N = 7;
    List<int> []g = new List<int>[N];
    for (int i = 0; i < g.Length; i++)
      g[i] = new List<int>();
 
    edge(1, 2, g);
    edge(1, 3, g);
    edge(2, 4, g);
    edge(2, 5, g);
    edge(5, 6, g);
    edge(5, 7, g);
    int[] res = pathSum(g);
 
    for (int i = 0; i < N; ++i) {
      Console.Write(res[i]+ " ");
    }
    Console.WriteLine();
  }
}
 
// This code is contributed by 29AjayKumar


Python3




## Python program for the above approach:
 
## Dfs computes dp, answer for each node
## with respect to its subtree it also
## computes size of each subtree
def dfs0(node, par, g, dp, size):
    ## Initialise given subtree with dp = 0
    ## [as there is no paths currently] and
    ## size 1, because tehre is only
    ## one node in subtree
    size[node] = 1
    dp[node] = 0
    for nebr in g[node]:
         
        ## For every neighbour of node
        ## which is not its parent
        ## 1. compute size and dp for
        ## nebr by dfs
        ## 2. update size and dp for node,
        ## based on nebr
        ## See explanation to understand
        ## the dp transition
        if (nebr != par):
            dfs0(nebr, node, g, dp, size)
            size[node] += size[nebr]
            dp[node] += (dp[nebr] + size[nebr])
 
## Rerooting the tree from 'from' to 'to'
def reroot(fr, to, dp, size):
    ## 'to' is no longer a child of 'from'
    dp[fr] -= size[to] + dp[to];
    size[fr] -= size[to];
      
    ## 'fr' is now a child of 'to'
    size[to] += size[fr];
    dp[to] += size[fr] + dp[fr];
  
def dfs1(node, par, g, dp, ans, size):
    ## Current dfs considers 'node' as root
    ## so currently dp[node]
    ## will be the answer
    ans[node] = dp[node];
  
    ## For all neighbours which are
    ## not parent of node
    for nebr in g[node]:
        if (par != nebr):
            ## Reroot the tree to 'nebr'
            reroot(node, nebr, dp, size);
  
            ## Compute ans for 'nebr'
            ## as a root of tree with dfs
            dfs1(nebr, node, g, dp, ans, size);
  
            ## reroot the tree back
            ## to 'node'
            reroot(nebr, node, dp, size);
 
## Creates a edge between a and b,
## given graph g
def edge(a, b, g):
    ## Convert into 0-based indexing
    a-=1
    b-=1
 
    ## Push b to adjacency list of a
    ## and vice versa because given
    ## tree is undirected
    g[a].append(b);
    g[b].append(a);
 
## Function to get the sum of paths
def pathSum(g, N):
    dp = [0]*N
    ans = [0]*N
    size = [0]*N
     
    ## Compute answer for each subtree
    ## with tree rooted at 0
    dfs0(0, -1, g, dp, size);
      
    ## Compute answer for each node
    ## as root of tree, rerooting
    dfs1(0, -1, g, dp, ans, size);
    return ans;
 
 
 
## Driver code
if __name__=='__main__':
 
    N = 7
     
    g = [0]*N
    for i in range(0, N):
        g[i] = []
     
    edge(1, 2, g)
    edge(1, 3, g)
    edge(2, 4, g)
    edge(2, 5, g)
    edge(5, 6, g)
    edge(5, 7, g)
    res = pathSum(g, N)
     
    for i in range(0, N):
        print(res[i], end=' ')
    print("")


Javascript




// JavaScript code to implement the approach
 
// Function to calculate the sum of paths for each node
function pathSum(g, N) {
  // Initialize dp and size arrays to store path sums
  // and subtree sizes for each node
  let dp = Array(N).fill(0),
    ans = Array(N).fill(0),
    size = Array(N).fill(0);
 
  // Function to compute dp and size for each node and its subtree
  function dfs0(node, par) {
    // Initialize given subtree with dp = 0 and size = 1
    dp[node] = 0;
    size[node] = 1;
    for (let nebr of g[node]) {
      // For every neighbor of node which is not its parent
      if (par !== nebr) {
        // Compute size and dp for neighbor by dfs
        dfs0(nebr, node);
        size[node] += size[nebr];
        dp[node] += size[nebr] + dp[nebr];
      }
    }
  }
 
  // Function to reroot the tree from 'from' to 'to'
  function reroot(from, to) {
    // 'to' is no longer a child of 'from'
    dp[from] -= size[to] + dp[to];
    size[from] -= size[to];
    // 'from' is now a child of 'to'
    size[to] += size[from];
    dp[to] += size[from] + dp[from];
  }
 
  // Function to compute ans for each node as root of the tree
  function dfs1(node, par) {
    // Current dfs considers 'node' as root
    // so currently dp[node] will be the answer
    ans[node] = dp[node];
 
    // For all neighbors which are not parent of node
    for (let nebr of g[node]) {
      if (par !== nebr) {
        // Reroot the tree to 'nebr'
        reroot(node, nebr);
        // Compute ans for 'nebr' as root of tree with dfs
        dfs1(nebr, node);
        // Reroot the tree back to 'node'
        reroot(nebr, node);
      }
    }
  }
 
  // Compute answer for each subtree with tree rooted at 0
  dfs0(0, -1);
  // Compute answer for each node as root of tree, rerooting
  dfs1(0, -1);
  return ans;
}
 
// Function to create an edge between a and b in graph g
function edge(a, b, g) {
  // Convert into 0-based indexing
  a--;
  b--;
  // Add b to adjacency list of a and vice versa
  // because the given tree is undirected
  g[a].push(b);
  g[b].push(a);
}
 
// Driver code
let N = 7;
let g = Array(N);
for (let i = 0; i < N; i++) g[i] = [];
 
edge(1, 2, g);
edge(1, 3, g);
edge(2, 4, g);
edge(2, 5, g);
edge(5, 6, g);
edge(5, 7, g);
    let res = pathSum(g, N);
    console.log(res);


Output

12 9 17 14 10 15 15 

Time Complexity: O(N)
Auxiliary Space: O(N)



Last Updated : 08 Feb, 2023
Like Article
Save Article
Previous
Next
Share your thoughts in the comments
Similar Reads