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’,
- Size of subtree ‘b’, size[b] and the sum of lengths of paths for subtree ‘b’ is given as dp[b].
- Consider all the paths which start at ‘a’ and end at some node in subtree of ‘b’.
- Every path can be broken down as follows: path(a, x) = edge(a, b) + path(b, x)
- 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’.
- 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:
- For each node do the following :
- Root the tree at this node, and find the dp as described above
- 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 subtreevoid 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 gvoid 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 pathsvector<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 codeint 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 approachimport 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 approachusing 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 subtreedef 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 gdef 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 pathsdef 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 codeif __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 subtreefunction 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 subtreevoid 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 gvoid 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 pathsvector<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 codeint 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 approachimport 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 approachusing 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 subtreedef 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 gdef 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 pathsdef 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 codeif __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 nodefunction 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 gfunction 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 codelet 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)
Feeling lost in the world of random DSA topics, wasting time without progress? It’s time for a change! Join our DSA course, where we’ll guide you on an exciting journey to master DSA efficiently and on schedule.
Ready to dive in? Explore our Free Demo Content and join our DSA course, trusted by over 100,000 neveropen!
Ready to dive in? Explore our Free Demo Content and join our DSA course, trusted by over 100,000 neveropen!
