Find distance between two nodes in the given Binary tree for Q queries | Python3 Program to find distance between two nodes using LCA ; lg2 ( MAX ) ; Array to store the level of each node ; Vector to store tree ; Pre - Processing to calculate values of lca [ ] [ ] , dist [ ] [ ] ; Using recursion formula to calculate the values of lca [ ] [ ] ; Storing the level of each node ; Using recursion formula to calculate the values of lca [ ] [ ] and dist [ ] [ ] ; Function to find the distance between given nodes u and v ; The node which is present farthest from the root node is taken as v . If u is farther from root node then swap the two ; Finding the ancestor of v which is at same level as u ; Adding distance of node v till its 2 ^ i - th ancestor ; If u is the ancestor of v then u is the LCA of u and v ; Finding the node closest to the root which is not the common ancestor of u and v i . e . a node x such that x is not the common ancestor of u and v but lca [ x ] [ 0 ] is ; Adding the distance of v and u to its 2 ^ i - th ancestor ; Adding the distance of u and v to its first ancestor ; Driver Code ; Number of nodes ; Add edges with their cost ; Initialising lca and dist values with - 1 and 0 respectively ; Perform DFS ; Query 1 : { 1 , 3 } ; Query 2 : { 2 , 3 } ; Query 3 : { 3 , 5 }
MAX = 1000 NEW_LINE lg = 10 NEW_LINE level = [ 0 for i in range ( MAX ) ] NEW_LINE lca = [ [ 0 for i in range ( lg ) ] for j in range ( MAX ) ] NEW_LINE dist = [ [ 0 for i in range ( lg ) ] for j in range ( MAX ) ] NEW_LINE graph = [ [ ] for i in range ( MAX ) ] NEW_LINE def addEdge ( u , v , cost ) : NEW_LINE INDENT global graph NEW_LINE graph [ u ] . append ( [ v , cost ] ) NEW_LINE graph [ v ] . append ( [ u , cost ] ) NEW_LINE DEDENT def dfs ( node , parent , h , cost ) : NEW_LINE INDENT lca [ node ] [ 0 ] = parent NEW_LINE level [ node ] = h NEW_LINE if ( parent != - 1 ) : NEW_LINE INDENT dist [ node ] [ 0 ] = cost NEW_LINE DEDENT for i in range ( 1 , lg ) : NEW_LINE INDENT if ( lca [ node ] [ i - 1 ] != - 1 ) : NEW_LINE INDENT lca [ node ] [ i ] = lca [ lca [ node ] [ i - 1 ] ] [ i - 1 ] NEW_LINE dist [ node ] [ i ] = ( dist [ node ] [ i - 1 ] + dist [ lca [ node ] [ i - 1 ] ] [ i - 1 ] ) NEW_LINE DEDENT DEDENT for i in graph [ node ] : NEW_LINE INDENT if ( i [ 0 ] == parent ) : NEW_LINE INDENT continue NEW_LINE DEDENT dfs ( i [ 0 ] , node , h + 1 , i [ 1 ] ) NEW_LINE DEDENT DEDENT def findDistance ( u , v ) : NEW_LINE INDENT ans = 0 NEW_LINE if ( level [ u ] > level [ v ] ) : NEW_LINE INDENT temp = u NEW_LINE u = v NEW_LINE v = temp NEW_LINE DEDENT i = lg - 1 NEW_LINE while ( i >= 0 ) : NEW_LINE INDENT if ( lca [ v ] [ i ] != - 1 and level [ lca [ v ] [ i ] ] >= level [ u ] ) : NEW_LINE INDENT ans += dist [ v ] [ i ] NEW_LINE v = lca [ v ] [ i ] NEW_LINE DEDENT i -= 1 NEW_LINE DEDENT if ( v == u ) : NEW_LINE INDENT print ( ans ) NEW_LINE DEDENT else : NEW_LINE INDENT i = lg - 1 NEW_LINE while ( i >= 0 ) : NEW_LINE INDENT if ( lca [ v ] [ i ] != lca [ u ] [ i ] ) : NEW_LINE INDENT ans += dist [ u ] [ i ] + dist [ v ] [ i ] NEW_LINE v = lca [ v ] [ i ] NEW_LINE u = lca [ u ] [ i ] NEW_LINE DEDENT i -= 1 NEW_LINE DEDENT ans += ( dist [ u ] [ 0 ] + dist [ v ] [ 0 ] ) NEW_LINE print ( ans ) NEW_LINE DEDENT DEDENT if __name__ == ' _ _ main _ _ ' : NEW_LINE INDENT n = 5 NEW_LINE addEdge ( 1 , 2 , 2 ) NEW_LINE addEdge ( 1 , 3 , 3 ) NEW_LINE addEdge ( 2 , 4 , 5 ) NEW_LINE addEdge ( 2 , 5 , 7 ) NEW_LINE for i in range ( 1 , n + 1 ) : NEW_LINE INDENT for j in range ( lg ) : NEW_LINE INDENT lca [ i ] [ j ] = - 1 NEW_LINE dist [ i ] [ j ] = 0 NEW_LINE DEDENT DEDENT dfs ( 1 , - 1 , 0 , 0 ) NEW_LINE findDistance ( 1 , 3 ) NEW_LINE findDistance ( 2 , 3 ) NEW_LINE findDistance ( 3 , 5 ) NEW_LINE DEDENT