Given a **weighted tree** containing **N** nodes, and two nodes **u** and **v**, the task is to find the count of nodes having **prime weight** on the simple path between **u and v (both inclusive)**.

**Examples:**

Input:

u = 3, v = 5

Output:2

Explanation:

Prime weight on path 3 to 5 is [11, 5]. Hence the answer is 2.

**Approach:** To solve th problem mentioned above, the idea is to use the basic concept when we find the LCA of two nodes.

- Precompute all the prime numbers till MAX using Seive method to check if a number is prime or not in O(1)
- Given two nodes u and v, we will make both nodes at the same level, by moving the greater level node move upwards. As we move up we will also check if the weight is prime or not.
- If
**v == u**then we will simply check the weight of current node and return the count. - If
**v is not equal to u**then we will move**both u and v upward by 1**till they are not same. - Now we will finally check the weight of the first ancestor of u or v and return the count.

Below is the implementation of the above approach:

`// C++ program Count prime weight ` `// nodes between two nodes in the given tree ` ` ` `#include <bits/stdc++.h> ` `using` `namespace` `std; ` `#define MAX 1000 ` ` ` `int` `weight[MAX]; ` `int` `level[MAX]; ` `int` `par[MAX]; ` `bool` `prime[MAX + 1]; ` `vector<` `int` `> graph[MAX]; ` ` ` `// Function to perform ` `// Sieve Of Eratosthenes for prime number ` `void` `SieveOfEratosthenes() ` `{ ` ` ` `// 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)); ` ` ` ` ` `for` `(` `int` `p = 2; p * p <= MAX; p++) { ` ` ` ` ` `// Check if prime[p] is not changed, ` ` ` `// then it is a prime ` ` ` `if` `(prime[p] == ` `true` `) { ` ` ` ` ` `// 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. ` ` ` `for` `(` `int` `i = p * p; i <= MAX; i += p) ` ` ` `prime[i] = ` `false` `; ` ` ` `} ` ` ` `} ` `} ` ` ` `// Function to perform dfs ` `void` `dfs(` `int` `node, ` `int` `parent, ` `int` `h) ` `{ ` ` ` `// Stores parent of each node ` ` ` `par[node] = parent; ` ` ` ` ` `// Stores level of each node from root ` ` ` `level[node] = h; ` ` ` ` ` `for` `(` `int` `child : graph[node]) { ` ` ` `if` `(child == parent) ` ` ` `continue` `; ` ` ` `dfs(child, node, h + 1); ` ` ` `} ` `} ` ` ` `// Function to perform prime ` `// number between the path ` `int` `findPrimeOnPath(` `int` `u, ` `int` `v) ` `{ ` ` ` `int` `count = 0; ` ` ` ` ` `// 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 ` ` ` `if` `(level[u] > level[v]) ` ` ` `swap(u, v); ` ` ` ` ` `int` `d = level[v] - level[u]; ` ` ` ` ` `// Find the ancestor of v ` ` ` `// which is at same level as u ` ` ` `while` `(d--) { ` ` ` ` ` `// If Weight is prime ` ` ` `// increment count ` ` ` `if` `(prime[weight[v]]) ` ` ` `count++; ` ` ` ` ` `v = par[v]; ` ` ` `} ` ` ` ` ` `// 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 ` ` ` `if` `(v == u) { ` ` ` `if` `(prime[weight[v]]) ` ` ` `count++; ` ` ` `return` `count; ` ` ` `} ` ` ` ` ` `// 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 ` ` ` `while` `(v != u) { ` ` ` `if` `(prime[weight[v]]) ` ` ` `count++; ` ` ` ` ` `if` `(prime[weight[u]]) ` ` ` `count++; ` ` ` ` ` `u = par[u]; ` ` ` `v = par[v]; ` ` ` `} ` ` ` `// If weight of first ancestor ` ` ` `// is prime ` ` ` `if` `(prime[weight[v]]) ` ` ` `count++; ` ` ` ` ` `return` `count; ` `} ` ` ` `// Driver code ` `int` `main() ` `{ ` ` ` `// Precompute all the prime ` ` ` `// numbers till MAX ` ` ` `SieveOfEratosthenes(); ` ` ` ` ` `// Weights of the node ` ` ` `weight[1] = 5; ` ` ` `weight[2] = 10; ` ` ` `weight[3] = 11; ` ` ` `weight[4] = 8; ` ` ` `weight[5] = 6; ` ` ` ` ` `// Edges of the tree ` ` ` `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; ` `} ` |

*chevron_right*

*filter_none*

**Output:**

2

**Time Complexity:** O(N)

Attention reader! Don’t stop learning now. Get hold of all the important DSA concepts with the **DSA Self Paced Course** at a student-friendly price and become industry ready.

## Recommended Posts:

- Count the nodes in the given tree whose weight is prime
- Count of Nodes which has Prime Digit sum weight in a Tree
- Count the nodes in the given tree whose weight is even
- Count the nodes in the given tree whose weight is even parity
- Count the nodes in the given tree whose weight is a power of two
- Count the nodes in the given tree whose sum of digits of weight is odd
- Count the nodes of the given tree whose weight has X as a factor
- Count nodes in the given tree whose weight is a fibonacci number
- Count the nodes in the given Tree whose weight is a Perfect Number
- Count the nodes in the given tree whose weight is a powerful number
- Count of nodes in a Binary Tree whose immediate children are co-prime
- Count of nodes in a Binary Tree whose child is its prime factors
- Count the nodes of the tree which make a pangram when concatenated with the sub-tree nodes
- Count the nodes whose weight is a perfect square
- Query to find the maximum and minimum weight between two nodes in the given tree using LCA.
- Print the nodes with a prime degree in given Prufer sequence of a Tree
- Print the nodes of the Binary Tree whose height is a Prime number
- Common nodes in the inorder sequence of a tree between given two nodes in O(1) space
- Count of Nodes whose both immediate children are its prime factors
- Count of Prime Nodes of a Singly Linked List

If you like GeeksforGeeks and would like to contribute, you can also write an article using contribute.geeksforgeeks.org or mail your article to contribute@geeksforgeeks.org. See your article appearing on the GeeksforGeeks main page and help other Geeks.

Please Improve this article if you find anything incorrect by clicking on the "Improve Article" button below.