Given a tree with N nodes and (N – 1) edges, where the nodes are assigned values from 1 to N, and the root node is 1. The task is to determine if there exists a subtree within the given tree that contains exactly K prime-numbered nodes.
Examples:
Input: N = 5, K = 2, edges[][] = {{1, 2}, {1, 3}, {2, 4}, {2, 5}} Output: 1 Explanation: Below image represents subtree containing 2 primes: 2 and 5

Input: N = 7, K = 5, edges = {{1, 2}, {1, 3}, {2, 4}, {2, 5}, {3, 6}, {3, 7}} Output: 0 Explanation: Below image shows that there is no subtree containing 5 primes:
-768.webp)
Approach: To solve the problem, follow the below idea:
The main idea is to use a Depth-First Search (DFS) to traverse the tree, counting the prime-numbered nodes in each subtree. The count is checked for equality with K, and if found, the answer is set to true, and here we can use the Sieve of Eratosthenes to efficiently identify prime numbers.
Step-by-step algorithm:
- Create an adjacency list adj for the tree and initialize a prime list prime from 0 to n, marking 0 and 1 as not prime using the Sieve of Eratosthenes.
- Define a DFS function (DFS(node, adj, prime, vis, ans, k)) that marks the current node as visited, counts prime nodes in its subtree, and increments ans if the count equals k.
- Recursively call the function for unvisited neighbors of the current node.
- Initialize a visited list vis and an answer list ans and call DFS for the root node (e.g., node 1).
- Check if the value in the ans list is greater than 0 and return true if yes, else return false.
Below is the implementation of the algorithm:
C++
#include <bits/stdc++.h>
using namespace std;
int dfs( int node, vector<vector< int > >& adj,
vector< bool >& prime, vector< int >& vis, int & ans,
int K)
{
if (ans > 0)
return 0;
vis[node] = 1;
int ele = 0;
if (prime[node])
ele += 1;
for ( auto j : adj[node]) {
if (vis[j] == 0) {
ele += dfs(j, adj, prime, vis, ans, K);
}
}
if (ele == K) {
ans += 1;
}
return ele;
}
bool hasKPrimeSubtree( int N, int K,
vector<vector< int > >& edges)
{
vector<vector< int > > adj(N + 1, vector< int >());
for ( auto j : edges) {
int u = j[0];
int v = j[1];
adj[u].push_back(v);
adj[v].push_back(u);
}
vector< bool > prime(N + 1, 1);
prime[0] = prime[1] = false ;
for ( int p = 2; p * p <= N; ++p) {
if (prime[p] == true ) {
for ( int i = p * p; i <= N; i += p)
prime[i] = false ;
}
}
vector< int > vis(N + 1, 0);
int ans = 0;
int ele = dfs(1, adj, prime, vis, ans, K);
return (ans > 0);
}
int main()
{
int N = 3;
int K = 1;
vector<vector< int > > edges = { { 1, 3 }, { 1, 2 } };
cout << hasKPrimeSubtree(N, K, edges);
return 0;
}
|
Java
import java.util.ArrayList;
import java.util.Arrays;
public class Main {
static int dfs( int node, ArrayList<ArrayList<Integer>> adj,
boolean [] prime, int [] vis, int [] ans, int K) {
if (ans[ 0 ] > 0 )
return 0 ;
vis[node] = 1 ;
int ele = 0 ;
if (prime[node])
ele += 1 ;
for ( int j : adj.get(node)) {
if (vis[j] == 0 ) {
ele += dfs(j, adj, prime, vis, ans, K);
}
}
if (ele == K) {
ans[ 0 ] += 1 ;
}
return ele;
}
static boolean hasKPrimeSubtree( int N, int K, ArrayList<ArrayList<Integer>> edges) {
ArrayList<ArrayList<Integer>> adj = new ArrayList<>();
for ( int i = 0 ; i <= N; i++) {
adj.add( new ArrayList<>());
}
for (ArrayList<Integer> j : edges) {
int u = j.get( 0 );
int v = j.get( 1 );
adj.get(u).add(v);
adj.get(v).add(u);
}
boolean [] prime = new boolean [N + 1 ];
Arrays.fill(prime, true );
prime[ 0 ] = prime[ 1 ] = false ;
for ( int p = 2 ; p * p <= N; ++p) {
if (prime[p]) {
for ( int i = p * p; i <= N; i += p)
prime[i] = false ;
}
}
int [] vis = new int [N + 1 ];
int [] ans = { 0 };
int ele = dfs( 1 , adj, prime, vis, ans, K);
return ans[ 0 ] > 0 ;
}
public static void main(String[] args) {
int N = 3 ;
int K = 1 ;
ArrayList<ArrayList<Integer>> edges = new ArrayList<>();
edges.add( new ArrayList<>(Arrays.asList( 1 , 3 )));
edges.add( new ArrayList<>(Arrays.asList( 1 , 2 )));
System.out.println(hasKPrimeSubtree(N, K, edges));
}
}
|
C#
using System;
using System.Collections.Generic;
class GFG
{
static int DFS( int node, List<List< int >> adj, List< bool > prime, List< int > vis, ref int ans, int K)
{
if (ans > 0)
return 0;
vis[node] = 1;
int ele = 0;
if (prime[node])
ele += 1;
foreach ( var j in adj[node])
{
if (vis[j] == 0)
{
ele += DFS(j, adj, prime, vis, ref ans, K);
}
}
if (ele == K)
{
ans += 1;
}
return ele;
}
static bool HasKPrimeSubtree( int N, int K, List<List< int >> edges)
{
List<List< int >> adj = new List<List< int >>();
for ( int i = 0; i <= N; i++)
{
adj.Add( new List< int >());
}
foreach ( var j in edges)
{
int u = j[0];
int v = j[1];
adj[u].Add(v);
adj[v].Add(u);
}
List< bool > prime = new List< bool >();
for ( int i = 0; i <= N; i++)
{
prime.Add( true );
}
prime[0] = prime[1] = false ;
for ( int p = 2; p * p <= N; ++p)
{
if (prime[p] == true )
{
for ( int i = p * p; i <= N; i += p)
prime[i] = false ;
}
}
List< int > vis = new List< int >();
for ( int i = 0; i <= N; i++)
{
vis.Add(0);
}
int ans = 0;
DFS(1, adj, prime, vis, ref ans, K);
return (ans > 0);
}
static void Main()
{
int N = 3;
int K = 1;
List<List< int >> edges = new List<List< int >> { new List< int > { 1, 3 }, new List< int > { 1, 2 } };
Console.WriteLine(HasKPrimeSubtree(N, K, edges));
}
}
|
Javascript
function dfs(node, adj, prime, vis, ans, K) {
if (ans[0] > 0) return 0;
vis[node] = 1;
let ele = 0;
if (prime[node]) ele += 1;
for (let j of adj[node]) {
if (vis[j] === 0) {
ele += dfs(j, adj, prime, vis, ans, K);
}
}
if (ele === K) {
ans[0] += 1;
}
return ele;
}
function hasKPrimeSubtree(N, K, edges) {
const adj = Array.from({ length: N + 1 }, () => []);
for (let j of edges) {
const u = j[0];
const v = j[1];
adj[u].push(v);
adj[v].push(u);
}
const prime = new Array(N + 1).fill( true );
prime[0] = prime[1] = false ;
for (let p = 2; p * p <= N; ++p) {
if (prime[p]) {
for (let i = p * p; i <= N; i += p) prime[i] = false ;
}
}
const vis = new Array(N + 1).fill(0);
const ans = [0];
const ele = dfs(1, adj, prime, vis, ans, K);
return ans[0] > 0;
}
const N = 3;
const K = 1;
const edges = [[1, 3], [1, 2]];
console.log(hasKPrimeSubtree(N, K, edges));
|
Python3
from collections import defaultdict
def dfs(node, adj, prime, vis, ans, K):
if ans[ 0 ] > 0 :
return 0
vis[node] = 1
ele = 0
if prime[node]:
ele + = 1
for j in adj[node]:
if vis[j] = = 0 :
ele + = dfs(j, adj, prime, vis, ans, K)
if ele = = K:
ans[ 0 ] + = 1
return ele
def hasKPrimeSubtree(N, K, edges):
adj = defaultdict( list )
for j in edges:
u = j[ 0 ]
v = j[ 1 ]
adj[u].append(v)
adj[v].append(u)
prime = [ True for _ in range (N + 1 )]
prime[ 0 ] = prime[ 1 ] = False
p = 2
while p * p < = N:
if prime[p] = = True :
for i in range (p * p, N + 1 , p):
prime[i] = False
p + = 1
vis = [ 0 for _ in range (N + 1 )]
ans = [ 0 ]
ele = dfs( 1 , adj, prime, vis, ans, K)
return (ans[ 0 ] > 0 )
if __name__ = = "__main__" :
N = 3
K = 1
edges = [[ 1 , 3 ], [ 1 , 2 ]]
print (hasKPrimeSubtree(N, K, edges))
|
Time Complexity: O(N * log(log(N))), where N is the number of nodes in the tree. Auxiliary Space: O(N)
|