Sum of lengths of all paths possible in a given tree | Python3 implementation of the approach ; Number of vertices ; Adjacency list representation of the tree ; Array that stores the subtree size ; Array to mark all the vertices which are visited ; Utility function to create an edge between two vertices ; Add a to b 's list ; Add b to a 's list ; Function to calculate the subtree size ; Mark visited ; For every adjacent node ; If not already visited ; Recursive call for the child ; Function to calculate the contribution of each edge ; Mark current node as visited ; For every adjacent node ; If it is not already visisted ; Function to return the required sum ; First pass of the dfs ; Second pass ; Driver code
sz = 10 ** 5 NEW_LINE n = 5 NEW_LINE an = 0 NEW_LINE tree = [ [ ] for i in range ( sz ) ] NEW_LINE subtree_size = [ 0 ] * sz NEW_LINE vis = [ 0 ] * sz NEW_LINE def addEdge ( a , b ) : NEW_LINE INDENT tree [ a ] . append ( b ) NEW_LINE tree [ b ] . append ( a ) NEW_LINE DEDENT def dfs ( node ) : NEW_LINE INDENT leaf = True NEW_LINE vis [ node ] = 1 NEW_LINE for child in tree [ node ] : NEW_LINE INDENT if ( vis [ child ] == 0 ) : NEW_LINE INDENT leaf = False NEW_LINE dfs ( child ) NEW_LINE subtree_size [ node ] += subtree_size [ child ] NEW_LINE DEDENT DEDENT if leaf : NEW_LINE INDENT subtree_size [ node ] = 1 NEW_LINE DEDENT DEDENT def contribution ( node , ans ) : NEW_LINE INDENT global an NEW_LINE vis [ node ] = 1 NEW_LINE for child in tree [ node ] : NEW_LINE INDENT if ( vis [ child ] == 0 ) : NEW_LINE INDENT an += ( subtree_size [ child ] * ( n - subtree_size [ child ] ) ) NEW_LINE contribution ( child , ans ) NEW_LINE DEDENT DEDENT DEDENT def getSum ( ) : NEW_LINE INDENT for i in range ( sz ) : NEW_LINE INDENT vis [ i ] = 0 NEW_LINE DEDENT dfs ( 0 ) NEW_LINE ans = 0 NEW_LINE for i in range ( sz ) : NEW_LINE INDENT vis [ i ] = 0 NEW_LINE DEDENT contribution ( 0 , ans ) NEW_LINE return an NEW_LINE DEDENT n = 5 NEW_LINE addEdge ( 0 , 1 ) NEW_LINE addEdge ( 0 , 2 ) NEW_LINE addEdge ( 1 , 3 ) NEW_LINE addEdge ( 1 , 4 ) NEW_LINE print ( getSum ( ) ) NEW_LINE