Skip to content
Related Articles

Related Articles

Improve Article
Save Article
Like Article

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

  • Last Updated : 23 Jan, 2022

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 tehre 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 tehre 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 tehre 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

 
 

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

 
 

Output
12 9 17 14 10 15 15 

 

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

 


My Personal Notes arrow_drop_up
Recommended Articles
Page :

Start Your Coding Journey Now!