text: Queries to find the count of shortest paths in a Tree that contains a given edge | Python3 implementation for the above approach ; Adjacency list to represent the tree ; Number of vertices ; Mark visited / unvisited vertices ; Stores the subtree size of the corresponding nodes ; Function to create an edge between two vertices ; Add a to b 's list ; Add b to a 's list ; Function to perform DFS ; Mark the vertex visited ; Include the node in the subtree ; Traverse all its children ; Function to print the required number of paths ; Driver Code ; Number of vertices ; Calling modified dfs function ; Count pairs of vertices in the tree
code: sz = 100000 NEW_LINE tree = [ [ ] for i in range ( sz ) ] NEW_LINE n = 0 NEW_LINE vis = [ False ] * sz NEW_LINE subtreeSize = [ 0 for i in range ( sz ) ] NEW_LINE def addEdge ( a , b ) : NEW_LINE INDENT global tree NEW_LINE tree [ a ] . append ( b ) NEW_LINE tree [ b ] . append ( a ) NEW_LINE DEDENT def dfs ( x ) : NEW_LINE INDENT global vis NEW_LINE global subtreeSize NEW_LINE global tree NEW_LINE vis [ x ] = True NEW_LINE subtreeSize [ x ] = 1 NEW_LINE for i in tree [ x ] : NEW_LINE INDENT if ( vis [ i ] == False ) : NEW_LINE INDENT dfs ( i ) NEW_LINE subtreeSize [ x ] += subtreeSize [ i ] NEW_LINE DEDENT DEDENT DEDENT def countPairs ( a , b ) : NEW_LINE INDENT global subtreeSize NEW_LINE sub = min ( subtreeSize [ a ] , subtreeSize [ b ] ) NEW_LINE print ( sub * ( n - sub ) ) NEW_LINE DEDENT if __name__ == ' _ _ main _ _ ' : NEW_LINE INDENT n = 6 NEW_LINE addEdge ( 0 , 1 ) NEW_LINE addEdge ( 0 , 2 ) NEW_LINE addEdge ( 1 , 3 ) NEW_LINE addEdge ( 3 , 4 ) NEW_LINE addEdge ( 3 , 5 ) NEW_LINE dfs ( 0 ) NEW_LINE countPairs ( 1 , 3 ) NEW_LINE countPairs ( 0 , 2 ) NEW_LINE DEDENT