Given a tree with N nodes, the task is to find the sum of lengths of all the paths. Path length for two nodes in the tree is the number of edges on the path and for two adjacent nodes in the tree, the length of the path is 1.
Examples:
Input:
0
/ \
1 2
/ \
3 4
Output: 18
Paths of length 1 = (0, 1), (0, 2), (1, 3), (1, 4) = 4
Paths of length 2 = (0, 3), (0, 4), (1, 2), (3, 4) = 4
Paths of length 3 = (3, 2), (4, 2) = 2
The sum of lengths of all paths =
(4 * 1) + (4 * 2) + (2 * 3) = 18
Input:
0
/
1
/
2
Output: 4
Naive approach: Check all possible paths and then add them to compute the final result. The complexity of this approach will be O(n2).
Efficient approach: It can be noted that each edge in a tree is a bridge. Hence that edge is going to be present in every path possible between the two subtrees that the edge connects.
For example, the edge (1 – 0) is present in every path possible between {1, 3, 4} and {0, 2}, (1 – 0) is used for 6 times that is size of the subtree {1, 3, 4} multiplied by the size of the subtree {0, 2}. So for each edge, compute how many times that edge is going to be considered for the paths going over it. DFS can be used to store the size of the subtree and the contribution of all edges can be computed with another dfs. The complexity of this approach will be O(n).
Below is the implementation of the above approach:
C++
// C++ implementation of the approach#include <bits/stdc++.h>using namespace std;const int sz = 1e5;// Number of verticesint n;// Adjacency list representation// of the treevector<int> tree[sz];// Array that stores the subtree sizeint subtree_size[sz];// Array to mark all the// vertices which are visitedint vis[sz];// Utility function to create an// edge between two verticesvoid addEdge(int a, int b){ // Add a to b's list tree[a].push_back(b); // Add b to a's list tree[b].push_back(a);}// Function to calculate the subtree sizeint dfs(int node){ // Mark visited vis[node] = 1; subtree_size[node] = 1; // For every adjacent node for (auto child : tree[node]) { // If not already visited if (!vis[child]) { // Recursive call for the child subtree_size[node] += dfs(child); } } return subtree_size[node];}// Function to calculate the// contribution of each edgevoid contribution(int node, int& ans){ // Mark current node as visited vis[node] = true; // For every adjacent node for (int child : tree[node]) { // If it is not already visited if (!vis[child]) { ans += (subtree_size[child] * (n - subtree_size[child])); contribution(child, ans); } }}// Function to return the required sumint getSum(){ // First pass of the dfs memset(vis, 0, sizeof(vis)); dfs(0); // Second pass int ans = 0; memset(vis, 0, sizeof(vis)); contribution(0, ans); return ans;}// Driver codeint main(){ n = 5; addEdge(0, 1); addEdge(0, 2); addEdge(1, 3); addEdge(1, 4); cout << getSum(); return 0;} |
Java
// Java implementation of the approach import java.util.*;@SuppressWarnings("unchecked") class GFG{ static int sz = 100005; // Number of vertices static int n; // Adjacency list representation // of the tree static ArrayList []tree = new ArrayList[sz]; // Array that stores the subtree size static int []subtree_size = new int[sz]; // Array to mark all the // vertices which are visited static int []vis = new int[sz]; // Utility function to create an // edge between two vertices static void AddEdge(int a, int b) { // Add a to b's list tree[a].add(b); // Add b to a's list tree[b].add(a); } // Function to calculate the subtree size static int dfs(int node) { // Mark visited vis[node] = 1; subtree_size[node] = 1; // For every adjacent node for(int child : (ArrayList<Integer>)tree[node]) { // If not already visited if (vis[child] == 0) { // Recursive call for the child subtree_size[node] += dfs(child); } } return subtree_size[node]; } // Function to calculate the // contribution of each edge static int contribution(int node, int ans) { // Mark current node as visited vis[node] = 1; // For every adjacent node for(int child : (ArrayList<Integer>)tree[node]) { // If it is not already visited if (vis[child] == 0) { ans += (subtree_size[child] * (n - subtree_size[child])); ans = contribution(child, ans); } } return ans;} // Function to return the required sum static int getSum() { // First pass of the dfs Arrays.fill(vis, 0); dfs(0); // Second pass int ans = 0; Arrays.fill(vis, 0); ans = contribution(0, ans); return ans; } // Driver code public static void main(String []args){ n = 5; for(int i = 0; i < sz; i++) { tree[i] = new ArrayList(); } AddEdge(0, 1); AddEdge(0, 2); AddEdge(1, 3); AddEdge(1, 4); System.out.println(getSum()); } }// This code is contributed by pratham76 |
Python3
# Python3 implementation of the approachsz = 10**5# Number of verticesn = 5an = 0# Adjacency list representation# of the treetree = [[] for i in range(sz)]# Array that stores the subtree sizesubtree_size = [0] * sz# Array to mark all the# vertices which are visitedvis = [0] * sz# Utility function to create an# edge between two verticesdef addEdge(a, b): # Add a to b's list tree[a].append(b) # Add b to a's list tree[b].append(a)# Function to calculate the subtree sizedef dfs(node): leaf = True # Mark visited vis[node] = 1 # For every adjacent node for child in tree[node]: # If not already visited if (vis[child] == 0): leaf = False dfs(child) # Recursive call for the child subtree_size[node] += subtree_size[child] if leaf: subtree_size[node] = 1# Function to calculate the# contribution of each edgedef contribution(node,ans): global an # Mark current node as visited vis[node] = 1 # For every adjacent node for child in tree[node]: # If it is not already visited if (vis[child] == 0): an += (subtree_size[child] * (n - subtree_size[child])) contribution(child, ans)# Function to return the required sumdef getSum(): # First pass of the dfs for i in range(sz): vis[i] = 0 dfs(0) # Second pass ans = 0 for i in range(sz): vis[i] = 0 contribution(0, ans) return an# Driver coden = 5addEdge(0, 1)addEdge(0, 2)addEdge(1, 3)addEdge(1, 4)print(getSum())# This code is contributed by Mohit Kumar |
C#
// C# implementation of the approach using System;using System.Collections;using System.Collections.Generic;class GFG{static int sz = 100005; // Number of vertices static int n; // Adjacency list representation // of the tree static ArrayList []tree = new ArrayList[sz]; // Array that stores the subtree size static int []subtree_size = new int[sz]; // Array to mark all the // vertices which are visited static int []vis = new int[sz]; // Utility function to create an // edge between two vertices static void addEdge(int a, int b) { // Add a to b's list tree[a].Add(b); // Add b to a's list tree[b].Add(a); } // Function to calculate the subtree size static int dfs(int node) { // Mark visited vis[node] = 1; subtree_size[node] = 1; // For every adjacent node foreach(int child in tree[node]) { // If not already visited if (vis[child] == 0) { // Recursive call for the child subtree_size[node] += dfs(child); } } return subtree_size[node]; } // Function to calculate the // contribution of each edge static void contribution(int node, ref int ans) { // Mark current node as visited vis[node] = 1; // For every adjacent node foreach(int child in tree[node]) { // If it is not already visited if (vis[child] == 0) { ans += (subtree_size[child] * (n - subtree_size[child])); contribution(child, ref ans); } } } // Function to return the required sum static int getSum() { // First pass of the dfs Array.Fill(vis, 0); dfs(0); // Second pass int ans = 0; Array.Fill(vis, 0); contribution(0, ref ans); return ans; } // Driver code public static void Main(){ n = 5; for(int i = 0; i < sz; i++) { tree[i] = new ArrayList(); } addEdge(0, 1); addEdge(0, 2); addEdge(1, 3); addEdge(1, 4); Console.Write(getSum()); } }// This code is contributed by rutvik_56 |
Javascript
<script>// Javascript implementation of the approachlet sz = 100005;// Number of verticeslet n;// Adjacency list representation// of the treelet tree = new Array(sz);// Array that stores the subtree sizelet subtree_size = new Array(sz);// Array to mark all the// vertices which are visitedlet vis = new Array(sz);// Utility function to create an// edge between two verticesfunction AddEdge(a,b){ // Add a to b's list tree[a].push(b); // Add b to a's list tree[b].push(a);}// Function to calculate the subtree sizefunction dfs(node){ // Mark visited vis[node] = 1; subtree_size[node] = 1; // For every adjacent node for(let child=0;child<tree[node].length;child++) { // If not already visited if (vis[tree[node][child]] == 0) { // Recursive call for the child subtree_size[node] += dfs(tree[node][child]); } } return subtree_size[node];}// Function to calculate the// contribution of each edgefunction contribution(node,ans){ // Mark current node as visited vis[node] = 1; // For every adjacent node for(let child=0;child<tree[node].length;child++) { // If it is not already visited if (vis[tree[node][child]] == 0) { ans += (subtree_size[tree[node][child]] * (n - subtree_size[tree[node][child]])); ans = contribution(tree[node][child], ans); } } return ans;}// Function to return the required sumfunction getSum(){ // First pass of the dfs for(let i=0;i<vis.length;i++) { vis[i]=0; } dfs(0); // Second pass let ans = 0; for(let i=0;i<vis.length;i++) { vis[i]=0; } ans = contribution(0, ans); return ans;}// Driver coden = 5; for(let i = 0; i < sz; i++){ tree[i] = [];}AddEdge(0, 1);AddEdge(0, 2);AddEdge(1, 3);AddEdge(1, 4);document.write(getSum());// This code is contributed by patel2127</script> |
Output:
18
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!
