Count nodes with prime weight in a Directed Acyclic Graph

Given a directed acyclic graph (DAG) with N nodes with integer representations 0 to N-1 and an integer array arr[], where arr[i] is the parent of node i. The root of the DAG is node i if arr[i] is 0. The weight of a node is calculated as the sum of the node’s value and the number of parent nodes it has, the task is to determine the number of nodes, excluding the root node, whose weight is a prime number.

Examples:

Input: N = 4, arr[] = {0, 1, 1, 2}
Output: 1
Explanation: The hierarchy is as follows

      (Root)
      Node 1
       /   \
 Node 2     Node 3  
     /
Node 4

Weight = Node + number of Parent Nodes
Weight of Node 1 = not considered.
Weight of Node 2 = 2+1 = 3 (prime)
Weight of Node 3 = 3+1 = 4 (not prime)
Weight of Node 4 = 4+2 = 6 (not prime)
Therefore, only Node 1 is prime

Input: N = 3, arr[] = {2, 3, 0}
Output: 2
Explanation: The hierarchy is as follows

      (Root)
      Node 3
       /   
 Node 2     
     /
Node 1

Weight of Node 3 = not considered. 
Weight of Node 2 = 2+1 = 3 (prime) 
Weight of Node 1 = 1+2 = 3 (prime)
Node 1 and 2 are both prime

Approach: To solve the problem follow the below idea:

The minimal distance from the Root node can be stored using DFS, and the number can be checked if it is prime or not using the Sieve of Eratosthenes.

Steps to solve the above problem:

• First, we will calculate all the prime numbers up to 10⁵.
• Use DFS (depth-first search) to determine the minimum distance from the Root node.
• Add the value of the current node and the minimum distance to obtain the weight of the node.
• Use the Sieve of Eratosthenes to determine whether or not that weight is a prime number.

Below is the code for the above approach:

C++

#include <algorithm>
#include <cstring>
#include <iostream>
using namespace std;

const int N = 100010;
bool b[N];
int dp[N];

void precompute()
{
    memset(b, true, sizeof(b));
    b[0] = false;
    b[1] = false;
    for (int i = 2; i < N; i++) {
        for (int j = i + i; j < N; j += i) {
            b[j] = false;
        }
    }
}

int solve(int i, int arr[])
{
    if (arr[i] == 0)
        return 0;
    if (dp[i] != -1)
        return dp[i];
    dp[i] = 1 + solve(arr[i] - 1, arr);
    return dp[i];
}

int primenode(int arr[], int n)
{
    precompute();
    memset(dp, -1, sizeof(dp));
    for (int i = 0; i < n; i++) {
        if (arr[i] == 0) {
            dp[i] = 0;
        }
    }
    for (int i = 0; i < n; i++) {
        if (dp[i] != -1)
            continue;
        dp[i] = solve(i, arr);
    }
    int cnt = 0;
    for (int i = 0; i < n; i++) {
        if (dp[i] == 0)
            continue;
        if (b[dp[i] + 1 + i]) {
            cnt++;
        }
    }
    return cnt;
}

int main()
{
    int arr[] = { 0, 1, 1, 2 };
    int n = sizeof(arr) / sizeof(arr[0]);
    int numPrimes = primenode(arr, n);
    cout << numPrimes << endl;
    return 0;
}

Java

import java.util.*;

class GFG {

    static int N = 100010;
    static boolean b[];

    // Function to precompute prime numbers
    // up to N
    static void precompute()
    {
        b = new boolean[N];
        Arrays.fill(b, true);
        b[0] = false;
        b[1] = false;
        for (int i = 2; i < 100010; i++) {
            for (int j = i + i; j < 100010; j = j + i) {
                b[j] = false;
            }
        }
    }

    // Function to recursively calculate
    // the number of subordinates
    // for each node
    static int solve(int i, int dp[], int arr[])
    {
        if (arr[i] == 0)
            return 0;
        if (dp[i] != -1)
            return dp[i];
        dp[i] = 1 + solve(arr[i] - 1, dp, arr);
        return dp[i];
    }

    // Function to calculate the
    // number of prime nodes
    static int primenode(int arr[], int n)
    {
        precompute();
        int dp[] = new int[n];
        Arrays.fill(dp, -1);
        for (int i = 0; i < n; i++) {
            if (arr[i] == 0) {
                dp[i] = 0;
            }
        }
        for (int i = 0; i < n; i++) {
            if (dp[i] != -1)
                continue;
            dp[i] = solve(i, dp, arr);
        }
        int cnt = 0;
        for (int i = 0; i < n; i++) {
            if (dp[i] == 0)
                continue;
            if (b[dp[i] + 1 + i]) {
                cnt++;
            }
        }
        return cnt;
    }

    // Main function to test
    // the implementation
    public static void main(String[] args)
    {
        int[] arr = { 0, 1, 1, 2 };
        int n = arr.length;
        int numPrimes = primenode(arr, n);
        System.out.println(numPrimes);
    }
}

Python3

# Python3 code for the approach

# Define a boolean array to store prime numbers
b = [True] * 100010

# Define a function to precompute prime numbers


def precompute():
    b[0] = False
    b[1] = False
    for i in range(2, 100010):
        for j in range(i + i, 100010, i):
            b[j] = False

# Define a function to calculate the length of the chain


def solve(i, arr, dp):
    if arr[i] == 0:
        return 0
    if dp[i] != -1:
        return dp[i]

    dp[i] = 1 + solve(arr[i] - 1, arr, dp)
    return dp[i]

# Define a function to calculate the number of prime nodes in the chain


def primenode(arr, n):
    # Precompute prime numbers
    precompute()
    # Initialize dp array with -1
    dp = [-1] * n
    for i in range(n):
        if arr[i] == 0:
            dp[i] = 0

    for i in range(n):
        if dp[i] != -1:
            continue

        dp[i] = solve(i, arr, dp)

    cnt = 0
    for i in range(n):
        if dp[i] == 0:
            continue
        if b[dp[i] + 1 + i]:
            cnt += 1

    return cnt

# func function to test
# the implementation


def func():
    arr = [0, 1, 1, 2]
    n = len(arr)
    numPrimes = primenode(arr, n)
    print(numPrimes)


if __name__ == '__main__':
            # Function call
    func()

C#

using System;

class GFG {

    static int N = 100010;
    static bool[] b;

    // Function to precompute prime numbers
    // up to N
    static void precompute()
    {
        b = new bool[N];
        Array.Fill(b, true);
        b[0] = false;
        b[1] = false;
        for (int i = 2; i < 100010; i++) {
            for (int j = i + i; j < 100010; j = j + i) {
                b[j] = false;
            }
        }
    }

    // Function to recursively calculate
    // the number of subordinates
    // for each node
    static int solve(int i, int[] dp, int[] arr)
    {
        if (arr[i] == 0)
            return 0;
        if (dp[i] != -1)
            return dp[i];
        dp[i] = 1 + solve(arr[i] - 1, dp, arr);
        return dp[i];
    }

    // Function to calculate the
    // number of prime nodes
    static int primenode(int[] arr, int n)
    {
        precompute();
        int[] dp = new int[n];
        Array.Fill(dp, -1);
        for (int i = 0; i < n; i++) {
            if (arr[i] == 0) {
                dp[i] = 0;
            }
        }
        for (int i = 0; i < n; i++) {
            if (dp[i] != -1)
                continue;
            dp[i] = solve(i, dp, arr);
        }
        int cnt = 0;
        for (int i = 0; i < n; i++) {
            if (dp[i] == 0)
                continue;
            if (b[dp[i] + 1 + i]) {
                cnt++;
            }
        }
        return cnt;
    }

    // Main function to test
    // the implementation
    public static void Main(string[] args)
    {
        int[] arr = { 0, 1, 1, 2 };
        int n = arr.Length;
        int numPrimes = primenode(arr, n);
        Console.WriteLine(numPrimes);
    }
}

Javascript

const N = 100010;
const b = new Array(N).fill(true);
const dp = new Array(N).fill(-1);

function precompute() {
  b[0] = false;
  b[1] = false;
  for (let i = 2; i < N; i++) {
    for (let j = i + i; j < N; j += i) {
      b[j] = false;
    }
  }
}

function solve(i, arr) {
  if (arr[i] == 0)
    return 0;
  if (dp[i] != -1)
    return dp[i];
  dp[i] = 1 + solve(arr[i] - 1, arr);
  return dp[i];
}

function primenode(arr, n) {
  precompute();
  for (let i = 0; i < n; i++) {
    if (arr[i] == 0) {
      dp[i] = 0;
    }
  }
  for (let i = 0; i < n; i++) {
    if (dp[i] != -1)
      continue;
    dp[i] = solve(i, arr);
  }
  let cnt = 0;
  for (let i = 0; i < n; i++) {
    if (dp[i] == 0)
      continue;
    if (b[dp[i] + 1 + i]) {
      cnt++;
    }
  }
  return cnt;
}

const arr = [0, 1, 1, 2];
const n = arr.length;
const numPrimes = primenode(arr, n);
console.log(numPrimes);

Output

1

Time Complexity: O(N)
Auxiliary Space: O(N)