text: Find the number of pair of Ideal nodes in a given tree | Python3 implementation of the approach ; Adjacency list ; bit : bit array i and j are starting and ending index INCLUSIVE ; bit : bit array n : size of bit array i is the index to be updated diff is ( new_val - old_val ) ; DFS function to find ideal pairs ; Function for initialisation ; Function to add an edge ; Function to find number of ideal pairs ; Find root of the tree ; Driver code ; Add edges ; Function call
code: N = 100005 NEW_LINE Ideal_pair = 0 NEW_LINE al = [ [ ] for i in range ( 100005 ) ] NEW_LINE bit = [ 0 for i in range ( N ) ] NEW_LINE root_node = [ 0 for i in range ( N ) ] NEW_LINE def bit_q ( i , j ) : NEW_LINE INDENT sum = 0 NEW_LINE while ( j > 0 ) : NEW_LINE INDENT sum += bit [ j ] NEW_LINE j -= ( j & ( j * - 1 ) ) NEW_LINE DEDENT i -= 1 NEW_LINE while ( i > 0 ) : NEW_LINE INDENT sum -= bit [ i ] NEW_LINE i -= ( i & ( i * - 1 ) ) NEW_LINE DEDENT return sum NEW_LINE DEDENT def bit_up ( i , diff ) : NEW_LINE INDENT while ( i <= n ) : NEW_LINE INDENT bit [ i ] += diff NEW_LINE i += i & - i NEW_LINE DEDENT DEDENT def dfs ( node , x ) : NEW_LINE INDENT Ideal_pair = x NEW_LINE Ideal_pair += bit_q ( max ( 1 , node - k ) , min ( n , node + k ) ) NEW_LINE bit_up ( node , 1 ) NEW_LINE for i in range ( len ( al [ node ] ) ) : NEW_LINE INDENT Ideal_pair = dfs ( al [ node ] [ i ] , Ideal_pair ) NEW_LINE DEDENT bit_up ( node , - 1 ) NEW_LINE return Ideal_pair NEW_LINE DEDENT def initialise ( ) : NEW_LINE INDENT Ideal_pair = 0 ; NEW_LINE for i in range ( n + 1 ) : NEW_LINE INDENT root_node [ i ] = True NEW_LINE bit [ i ] = 0 NEW_LINE DEDENT DEDENT def Add_Edge ( x , y ) : NEW_LINE INDENT al [ x ] . append ( y ) NEW_LINE root_node [ y ] = False NEW_LINE DEDENT def Idealpairs ( ) : NEW_LINE INDENT r = - 1 NEW_LINE for i in range ( 1 , n + 1 , 1 ) : NEW_LINE INDENT if ( root_node [ i ] ) : NEW_LINE INDENT r = i NEW_LINE break NEW_LINE DEDENT DEDENT Ideal_pair = dfs ( r , 0 ) NEW_LINE return Ideal_pair NEW_LINE DEDENT if __name__ == ' _ _ main _ _ ' : NEW_LINE INDENT n = 6 NEW_LINE k = 3 NEW_LINE initialise ( ) NEW_LINE Add_Edge ( 1 , 2 ) NEW_LINE Add_Edge ( 1 , 3 ) NEW_LINE Add_Edge ( 3 , 4 ) NEW_LINE Add_Edge ( 3 , 5 ) NEW_LINE Add_Edge ( 3 , 6 ) NEW_LINE print ( Idealpairs ( ) ) NEW_LINE DEDENT