text: Queries to find the count of shortest paths in a Tree that contains a given edge | C ++ 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: #include <bits/stdc++.h> NEW_LINE using namespace std ; const int sz = 1e5 ; vector < int > tree [ sz ] ; int n ; bool vis [ sz ] ; int subtreeSize [ sz ] ; void addEdge ( int a , int b ) { tree [ a ] . push_back ( b ) ; tree [ b ] . push_back ( a ) ; } void dfs ( int x ) { vis [ x ] = true ; subtreeSize [ x ] = 1 ; for ( auto i : tree [ x ] ) { if ( ! vis [ i ] ) { dfs ( i ) ; subtreeSize [ x ] += subtreeSize [ i ] ; } } } void countPairs ( int a , int b ) { int sub = min ( subtreeSize [ a ] , subtreeSize [ b ] ) ; cout << sub * ( n - sub ) << endl ; } int main ( ) { n = 6 ; addEdge ( 0 , 1 ) ; addEdge ( 0 , 2 ) ; addEdge ( 1 , 3 ) ; addEdge ( 3 , 4 ) ; addEdge ( 3 , 5 ) ; dfs ( 0 ) ; countPairs ( 1 , 3 ) ; countPairs ( 0 , 2 ) ; return 0 ; }