Subtree with exactly K primes

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:

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;

// Function to start iterating from node 1
int dfs(int node, vector<vector<int> >& adj,
        vector<bool>& prime, vector<int>& vis, int& ans,
        int K)
{

    if (ans > 0)
        return 0;

    // If node is visited
    vis[node] = 1;

    int ele = 0;

    // If node is prime
    if (prime[node])
        ele += 1;

    // Iterate further to its adjacent nodes
    for (auto j : adj[node]) {
        if (vis[j] == 0) {
            ele += dfs(j, adj, prime, vis, ans, K);
        }
    }

    // If total prime numbers are K
    if (ele == K) {
        ans += 1;
    }

    // Return ele
    return ele;
}

// Function to find whether tree have
// K prime subset
bool hasKPrimeSubtree(int N, int K,
                    vector<vector<int> >& edges)
{
    // Create a adjancy matrix
    vector<vector<int> > adj(N + 1, vector<int>());

    // Iterate in edges
    for (auto j : edges) {

        int u = j[0];
        int v = j[1];
        adj[u].push_back(v);
        adj[v].push_back(u);
    }

    // Create a prime vector
    vector<bool> prime(N + 1, 1);

    // As 0 and 1 are not prime number
    prime[0] = prime[1] = false;

    // Fill the prime vectors
    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 to keep a check on visited elements
    vector<int> vis(N + 1, 0);
    int ans = 0;

    // Function to start iterating in tree
    int ele = dfs(1, adj, prime, vis, ans, K);
    return (ans > 0);
}

// Driver code
int main()
{

    int N = 3;
    int K = 1;
    vector<vector<int> > edges = { { 1, 3 }, { 1, 2 } };

    // Function call
    cout << hasKPrimeSubtree(N, K, edges);
    return 0;
}

Java

import java.util.ArrayList;
import java.util.Arrays;

public class Main {

    // Function to start iterating from node 1
    static int dfs(int node, ArrayList<ArrayList<Integer>> adj,
                   boolean[] prime, int[] vis, int[] ans, int K) {

        if (ans[0] > 0)
            return 0;

        // If node is visited
        vis[node] = 1;

        int ele = 0;

        // If node is prime
        if (prime[node])
            ele += 1;

        // Iterate further to its adjacent nodes
        for (int j : adj.get(node)) {
            if (vis[j] == 0) {
                ele += dfs(j, adj, prime, vis, ans, K);
            }
        }

        // If total prime numbers are K
        if (ele == K) {
            ans[0] += 1;
        }

        // Return ele
        return ele;
    }

    // Function to find whether tree has K prime subset
    static boolean hasKPrimeSubtree(int N, int K, ArrayList<ArrayList<Integer>> edges) {
        // Create an adjacency matrix
        ArrayList<ArrayList<Integer>> adj = new ArrayList<>();
        for (int i = 0; i <= N; i++) {
            adj.add(new ArrayList<>());
        }

        // Iterate in edges
        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);
        }

        // Create a prime array
        boolean[] prime = new boolean[N + 1];
        Arrays.fill(prime, true);

        // As 0 and 1 are not prime numbers
        prime[0] = prime[1] = false;

        // Fill the prime array
        for (int p = 2; p * p <= N; ++p) {
            if (prime[p]) {
                for (int i = p * p; i <= N; i += p)
                    prime[i] = false;
            }
        }

        // Array to keep a check on visited elements
        int[] vis = new int[N + 1];
        int[] ans = { 0 };

        // Function to start iterating in the tree
        int ele = dfs(1, adj, prime, vis, ans, K);
        return ans[0] > 0;
    }

    // Driver code
    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)));

        // Function call
        System.out.println(hasKPrimeSubtree(N, K, edges));
    }
}

// This code is contributed by akshitaguprzj3

C#

using System;
using System.Collections.Generic;

class GFG
{
    // Function to start iterating from node 1
    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;

        // If node is visited
        vis[node] = 1;

        int ele = 0;

        // If node is prime
        if (prime[node])
            ele += 1;

        // Iterate further to its adjacent nodes
        foreach (var j in adj[node])
        {
            if (vis[j] == 0)
            {
                ele += DFS(j, adj, prime, vis, ref ans, K);
            }
        }

        // If total prime numbers are K
        if (ele == K)
        {
            ans += 1;
        }

        // Return ele
        return ele;
    }

    // Function to find whether tree has K prime subset
    static bool HasKPrimeSubtree(int N, int K, List<List<int>> edges)
    {
        // Create an adjacency matrix
        List<List<int>> adj = new List<List<int>>();
        for (int i = 0; i <= N; i++)
        {
            adj.Add(new List<int>());
        }

        // Iterate in edges
        foreach (var j in edges)
        {
            int u = j[0];
            int v = j[1];
            adj[u].Add(v);
            adj[v].Add(u);
        }

        // Create a prime vector
        List<bool> prime = new List<bool>();
        for (int i = 0; i <= N; i++)
        {
            prime.Add(true);
        }

        // As 0 and 1 are not prime numbers
        prime[0] = prime[1] = false;

        // Fill the prime vectors
        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 to keep a check on visited elements
        List<int> vis = new List<int>();
        for (int i = 0; i <= N; i++)
        {
            vis.Add(0);
        }
        int ans = 0;

        // Function to start iterating in tree
        DFS(1, adj, prime, vis, ref ans, K);
        return (ans > 0);
    }

    // Driver code
    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 } };

        // Function call
        Console.WriteLine(HasKPrimeSubtree(N, K, edges));
    }
}

Javascript

function dfs(node, adj, prime, vis, ans, K) {
    if (ans[0] > 0) return 0;

    // If node is visited
    vis[node] = 1;

    let ele = 0;

    // If node is prime
    if (prime[node]) ele += 1;

    // Iterate further to its adjacent nodes
    for (let j of adj[node]) {
        if (vis[j] === 0) {
            ele += dfs(j, adj, prime, vis, ans, K);
        }
    }

    // If total prime numbers are K
    if (ele === K) {
        ans[0] += 1;
    }

    // Return ele
    return ele;
}

function hasKPrimeSubtree(N, K, edges) {
    // Create an adjacency matrix
    const adj = Array.from({ length: N + 1 }, () => []);

    // Iterate in edges
    for (let j of edges) {
        const u = j[0];
        const v = j[1];
        adj[u].push(v);
        adj[v].push(u);
    }

    // Create a prime array
    const prime = new Array(N + 1).fill(true);

    // As 0 and 1 are not prime numbers
    prime[0] = prime[1] = false;

    // Fill the prime array
    for (let p = 2; p * p <= N; ++p) {
        if (prime[p]) {
            for (let i = p * p; i <= N; i += p) prime[i] = false;
        }
    }

    // Array to keep a check on visited elements
    const vis = new Array(N + 1).fill(0);
    const ans = [0];

    // Function to start iterating in the tree
    const ele = dfs(1, adj, prime, vis, ans, K);
    return ans[0] > 0;
}

// Driver code
const N = 3;
const K = 1;
const edges = [[1, 3], [1, 2]];

// Function call
console.log(hasKPrimeSubtree(N, K, edges));

Python3

from collections import defaultdict

# Function to start iterating from node 1
def dfs(node, adj, prime, vis, ans, K):
    if ans[0] > 0:
        return 0

    # If node is visited
    vis[node] = 1

    ele = 0

    # If node is prime
    if prime[node]:
        ele += 1

    # Iterate further to its adjacent nodes
    for j in adj[node]:
        if vis[j] == 0:
            ele += dfs(j, adj, prime, vis, ans, K)

    # If total prime numbers are K
    if ele == K:
        ans[0] += 1

    # Return ele
    return ele

# Function to find whether tree have K prime subset
def hasKPrimeSubtree(N, K, edges):
    # Create a adjacency matrix
    adj = defaultdict(list)

    # Iterate in edges
    for j in edges:
        u = j[0]
        v = j[1]
        adj[u].append(v)
        adj[v].append(u)

    # Create a prime list
    prime = [True for _ in range(N + 1)]

    # As 0 and 1 are not prime number
    prime[0] = prime[1] = False

    # Fill the prime list
    p = 2
    while p * p <= N:
        if prime[p] == True:
            for i in range(p * p, N + 1, p):
                prime[i] = False
        p += 1

    # List to keep a check on visited elements
    vis = [0 for _ in range(N + 1)]
    ans = [0]

    # Function to start iterating in tree
    ele = dfs(1, adj, prime, vis, ans, K)
    return (ans[0] > 0)

# Driver code
if __name__ == "__main__":
    N = 3
    K = 1
    edges = [[1, 3], [1, 2]]

    # Function call
    print(hasKPrimeSubtree(N, K, edges))

Output

1

Time Complexity: O(N * log(log(N))), where N is the number of nodes in the tree.
Auxiliary Space: O(N)