Count of all prime weight nodes between given nodes in the given Tree | Python3 program count prime weight nodes between two nodes in the given tree ; Function to perform Sieve Of Eratosthenes for prime number ; Initialize all entries of prime it as true . A value in prime [ i ] will finally be false if i is Not a prime , else true . memset ( prime , true , sizeof ( prime ) ) ; Check if prime [ p ] is not changed , then it is a prime ; Update all multiples of p greater than or equal to the square of it numbers which are multiple of p and are less than p ^ 2 are already been marked . ; Function to perform dfs ; Stores parent of each node ; Stores level of each node from root ; Function to perform prime number between the path ; 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 ; Find the ancestor of v which is at same level as u ; If Weight is prime increment count ; If u is the ancestor of v then u is the LCA of u and v Now check if weigh [ v ] is prime or not ; When v and u are on the same level but are in different subtree . Now move both u and v up by 1 till they are not same ; If weight of first ancestor is prime ; Driver code ; Precompute all the prime numbers till MAX ; Weights of the node ; Edges of the tree
MAX = 1000 NEW_LINE weight = [ 0 for i in range ( MAX ) ] NEW_LINE level = [ 0 for i in range ( MAX ) ] NEW_LINE par = [ 0 for i in range ( MAX ) ] NEW_LINE prime = [ True for i in range ( MAX + 1 ) ] NEW_LINE graph = [ [ ] for i in range ( MAX ) ] NEW_LINE def SieveOfEratosthenes ( ) : NEW_LINE INDENT for p in range ( 2 , MAX + 1 ) : NEW_LINE INDENT if p * p > MAX + 1 : NEW_LINE INDENT break NEW_LINE DEDENT if ( prime [ p ] == True ) : NEW_LINE INDENT for i in range ( p * p , MAX + 1 , p ) : NEW_LINE INDENT prime [ i ] = False NEW_LINE DEDENT DEDENT DEDENT DEDENT def dfs ( node , parent , h ) : NEW_LINE INDENT par [ node ] = parent NEW_LINE level [ node ] = h NEW_LINE for child in graph [ node ] : NEW_LINE INDENT if ( child == parent ) : NEW_LINE INDENT continue NEW_LINE DEDENT dfs ( child , node , h + 1 ) NEW_LINE DEDENT DEDENT def findPrimeOnPath ( u , v ) : NEW_LINE INDENT count = 0 NEW_LINE if ( level [ u ] > level [ v ] ) : NEW_LINE INDENT u , v = v , u NEW_LINE DEDENT d = level [ v ] - level [ u ] NEW_LINE while ( d ) : NEW_LINE INDENT if ( prime [ weight [ v ] ] ) : NEW_LINE INDENT count += 1 NEW_LINE DEDENT v = par [ v ] NEW_LINE d -= 1 NEW_LINE DEDENT if ( v == u ) : NEW_LINE INDENT if ( prime [ weight [ v ] ] ) : NEW_LINE INDENT count += 1 NEW_LINE DEDENT return count NEW_LINE DEDENT while ( v != u ) : NEW_LINE INDENT if ( prime [ weight [ v ] ] ) : NEW_LINE INDENT count += 1 NEW_LINE DEDENT if ( prime [ weight [ u ] ] ) : NEW_LINE INDENT count += 1 NEW_LINE DEDENT u = par [ u ] NEW_LINE v = par [ v ] NEW_LINE DEDENT if ( prime [ weight [ v ] ] ) : NEW_LINE INDENT count += 1 NEW_LINE DEDENT return count NEW_LINE DEDENT if __name__ == ' _ _ main _ _ ' : NEW_LINE INDENT SieveOfEratosthenes ( ) NEW_LINE weight [ 1 ] = 5 NEW_LINE weight [ 2 ] = 10 NEW_LINE weight [ 3 ] = 11 NEW_LINE weight [ 4 ] = 8 NEW_LINE weight [ 5 ] = 6 NEW_LINE graph [ 1 ] . append ( 2 ) NEW_LINE graph [ 2 ] . append ( 3 ) NEW_LINE graph [ 2 ] . append ( 4 ) NEW_LINE graph [ 1 ] . append ( 5 ) NEW_LINE dfs ( 1 , - 1 , 0 ) NEW_LINE u = 3 NEW_LINE v = 5 NEW_LINE print ( findPrimeOnPath ( u , v ) ) NEW_LINE DEDENT