Count of nodes having odd divisors in the given subtree for Q queries | Python3 implementation to count the number of nodes having odd number of divisors for each query ; Adjacency list for tree . ; Array for values and answer at ith node . ; Function to check whether N has odd divisors or not ; DFS function to pre - compute the answers ; Initialize the count ; Repeat for every child ; Increase the count if current node has odd number of divisors ; Driver Code ; Adjacency List ; Function call
import math NEW_LINE N = 100001 NEW_LINE adj = [ [ ] for i in range ( N ) ] NEW_LINE a = [ 0 for i in range ( N ) ] NEW_LINE ans = [ 0 for i in range ( N ) ] NEW_LINE def hasOddNumberOfDivisors ( n ) : NEW_LINE INDENT if ( math . sqrt ( n ) == int ( math . sqrt ( n ) ) ) : NEW_LINE INDENT return True NEW_LINE DEDENT return False NEW_LINE DEDENT def dfs ( node , parent ) : NEW_LINE INDENT count = 0 NEW_LINE for i in adj [ node ] : NEW_LINE INDENT if ( i != parent ) : NEW_LINE INDENT count += dfs ( i , node ) NEW_LINE DEDENT DEDENT if ( hasOddNumberOfDivisors ( a [ node ] ) ) : NEW_LINE INDENT count += 1 NEW_LINE DEDENT ans [ node ] = count NEW_LINE return count NEW_LINE DEDENT if __name__ == " _ _ main _ _ " : NEW_LINE INDENT n = 5 NEW_LINE i = 0 NEW_LINE q = [ 4 , 1 , 5 , 3 ] NEW_LINE adj [ 1 ] . append ( 2 ) NEW_LINE adj [ 2 ] . append ( 1 ) NEW_LINE adj [ 2 ] . append ( 3 ) NEW_LINE adj [ 3 ] . append ( 2 ) NEW_LINE adj [ 3 ] . append ( 4 ) NEW_LINE adj [ 4 ] . append ( 3 ) NEW_LINE adj [ 1 ] . append ( 5 ) NEW_LINE adj [ 5 ] . append ( 1 ) NEW_LINE a [ 1 ] = 4 NEW_LINE a [ 2 ] = 9 NEW_LINE a [ 3 ] = 14 NEW_LINE a [ 4 ] = 100 NEW_LINE a [ 5 ] = 5 NEW_LINE dfs ( 1 , - 1 ) NEW_LINE for i in range ( len ( q ) ) : NEW_LINE INDENT print ( ans [ q [ i ] ] , end = ' ▁ ' ) NEW_LINE DEDENT DEDENT