Count of all prime weight nodes between given nodes in the given Tree
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 the 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 the Sieve 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 the 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 the 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++
#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];
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;
}
|
Java
import java.util.*;
class GFG{
static final int MAX = 1000 ;
static int []weight =
new int [MAX];
static int []level =
new int [MAX];
static int []par =
new int [MAX];
static boolean []prime =
new boolean [MAX + 1 ];
static Vector<Integer>[] graph =
new Vector[MAX];
static void SieveOfEratosthenes()
{
for ( int i = 0 ;
i < prime.length; i++)
prime[i] = true ;
for ( int p = 2 ;
p * p <= MAX; p++)
{
if (prime[p] == true )
{
for ( int i = p * p;
i <= MAX; i += p)
prime[i] = false ;
}
}
}
static 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 );
}
}
static int findPrimeOnPath( int u,
int v)
{
int count = 0 ;
if (level[u] > level[v])
{
int temp = v;
v = u;
u = temp;
}
int d = level[v] - level[u];
while (d-- > 0 )
{
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;
}
public static void main(String[] args)
{
for ( int i = 0 ; i < graph.length; i++)
graph[i] = new Vector<Integer>();
SieveOfEratosthenes();
weight[ 1 ] = 5 ;
weight[ 2 ] = 10 ;
weight[ 3 ] = 11 ;
weight[ 4 ] = 8 ;
weight[ 5 ] = 6 ;
graph[ 1 ].add( 2 );
graph[ 2 ].add( 3 );
graph[ 2 ].add( 4 );
graph[ 1 ].add( 5 );
dfs( 1 , - 1 , 0 );
int u = 3 , v = 5 ;
System.out.print(findPrimeOnPath(u, v));
}
}
|
Python3
MAX = 1000
weight = [ 0 for i in range ( MAX )]
level = [ 0 for i in range ( MAX )]
par = [ 0 for i in range ( MAX )]
prime = [ True for i in range ( MAX + 1 )]
graph = [[] for i in range ( MAX )]
def SieveOfEratosthenes():
for p in range ( 2 , MAX + 1 ):
if p * p > MAX + 1 :
break
if (prime[p] = = True ):
for i in range (p * p, MAX + 1 , p):
prime[i] = False
def dfs(node, parent, h):
par[node] = parent
level[node] = h
for child in graph[node]:
if (child = = parent):
continue
dfs(child, node, h + 1 )
def findPrimeOnPath(u, v):
count = 0
if (level[u] > level[v]):
u, v = v, u
d = level[v] - level[u]
while (d):
if (prime[weight[v]]):
count + = 1
v = par[v]
d - = 1
if (v = = u):
if (prime[weight[v]]):
count + = 1
return count
while (v ! = u):
if (prime[weight[v]]):
count + = 1
if (prime[weight[u]]):
count + = 1
u = par[u]
v = par[v]
if (prime[weight[v]]):
count + = 1
return count
if __name__ = = '__main__' :
SieveOfEratosthenes()
weight[ 1 ] = 5
weight[ 2 ] = 10
weight[ 3 ] = 11
weight[ 4 ] = 8
weight[ 5 ] = 6
graph[ 1 ].append( 2 )
graph[ 2 ].append( 3 )
graph[ 2 ].append( 4 )
graph[ 1 ].append( 5 )
dfs( 1 , - 1 , 0 )
u = 3
v = 5
print (findPrimeOnPath(u, v))
|
C#
using System;
using System.Collections.Generic;
class GFG{
static readonly int MAX = 1000;
static int []weight = new int [MAX];
static int []level = new int [MAX];
static int []par = new int [MAX];
static bool []prime = new bool [MAX + 1];
static List< int >[] graph = new List< int >[MAX];
static void SieveOfEratosthenes()
{
for ( int i = 0;
i < prime.Length; i++)
prime[i] = true ;
for ( int p = 2;
p * p <= MAX; p++)
{
if (prime[p] == true )
{
for ( int i = p * p;
i <= MAX; i += p)
prime[i] = false ;
}
}
}
static void dfs( int node, int parent,
int h)
{
par[node] = parent;
level[node] = h;
foreach ( int child in graph[node])
{
if (child == parent)
continue ;
dfs(child, node, h + 1);
}
}
static int findPrimeOnPath( int u,
int v)
{
int count = 0;
if (level[u] > level[v])
{
int temp = v;
v = u;
u = temp;
}
int d = level[v] - level[u];
while (d-- > 0)
{
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;
}
public static void Main(String[] args)
{
for ( int i = 0; i < graph.Length; i++)
graph[i] = new List< int >();
SieveOfEratosthenes();
weight[1] = 5;
weight[2] = 10;
weight[3] = 11;
weight[4] = 8;
weight[5] = 6;
graph[1].Add(2);
graph[2].Add(3);
graph[2].Add(4);
graph[1].Add(5);
dfs(1, -1, 0);
int u = 3, v = 5;
Console.Write(findPrimeOnPath(u, v));
}
}
|
Javascript
<script>
var MAX = 1000;
var weight = Array(MAX);
var level = Array(MAX);
var par = Array(MAX);
var prime = Array(MAX+1).fill( true );
var graph = Array.from(Array(MAX), ()=>Array());
function SieveOfEratosthenes()
{
for ( var p = 2; p * p <= MAX; p++) {
if (prime[p] == true ) {
for ( var i = p * p; i <= MAX; i += p)
prime[i] = false ;
}
}
}
function dfs(node, parent, h)
{
par[node] = parent;
level[node] = h;
graph[node].forEach(child => {
if (child != parent)
dfs(child, node, h + 1);
});
}
function findPrimeOnPath(u, v)
{
var count = 0;
if (level[u] > level[v])
{
[u,v] = [v,u]
}
var 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;
}
SieveOfEratosthenes();
weight[1] = 5;
weight[2] = 10;
weight[3] = 11;
weight[4] = 8;
weight[5] = 6;
graph[1].push(2);
graph[2].push(3);
graph[2].push(4);
graph[1].push(5);
dfs(1, -1, 0);
var u = 3, v = 5;
document.write( findPrimeOnPath(u, v));
</script>
|
Complexity Analysis:
- Time Complexity: O(N).
In dfs, every node of the tree is processed once, and hence the complexity due to the dfs is O(N) if there are total N nodes in the tree. Also, for processing each node the SieveOfEratosthenes() function is used which has a complexity of O(sqrt(N)) too but since this function is executed only once, it does not affect the overall time complexity. Therefore, the time complexity is O(N).
- Auxiliary Space: O(N).
Extra space is used for the prime array, so the space complexity is O(N).
Last Updated :
03 Jun, 2021
Like Article
Save Article
Share your thoughts in the comments
Please Login to comment...