Count of all prime weight nodes between given nodes in the given Tree | C ++ 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 . ; 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
#include <bits/stdc++.h> NEW_LINE using namespace std ; #define MAX  1000 NEW_LINE int weight [ MAX ] ; int level [ MAX ] ; int par [ MAX ] ; bool prime [ MAX + 1 ] ; vector < int > graph [ MAX ] ; void SieveOfEratosthenes ( ) { memset ( prime , true , sizeof ( prime ) ) ; for ( int p = 2 ; p * p <= MAX ; p ++ ) { if ( prime [ p ] == true ) { for ( int i = p * p ; i <= MAX ; i += p ) prime [ i ] = false ; } } } void dfs ( int node , int parent , int h ) { par [ node ] = parent ; level [ node ] = h ; for ( int child : graph [ node ] ) { if ( child == parent ) continue ; dfs ( child , node , h + 1 ) ; } } int findPrimeOnPath ( int u , int v ) { int count = 0 ; if ( level [ u ] > level [ v ] ) swap ( u , v ) ; int d = level [ v ] - level [ u ] ; while ( d -- ) { if ( prime [ weight [ v ] ] ) count ++ ; v = par [ v ] ; } if ( v == u ) { if ( prime [ weight [ v ] ] ) count ++ ; return count ; } while ( v != u ) { if ( prime [ weight [ v ] ] ) count ++ ; if ( prime [ weight [ u ] ] ) count ++ ; u = par [ u ] ; v = par [ v ] ; } if ( prime [ weight [ v ] ] ) count ++ ; return count ; } int main ( ) { SieveOfEratosthenes ( ) ; weight [ 1 ] = 5 ; weight [ 2 ] = 10 ; weight [ 3 ] = 11 ; weight [ 4 ] = 8 ; weight [ 5 ] = 6 ; graph [ 1 ] . push_back ( 2 ) ; graph [ 2 ] . push_back ( 3 ) ; graph [ 2 ] . push_back ( 4 ) ; graph [ 1 ] . push_back ( 5 ) ; dfs ( 1 , -1 , 0 ) ; int u = 3 , v = 5 ; cout << findPrimeOnPath ( u , v ) << endl ; return 0 ; }