树中具有质数总和权重的节点数

给定加权树，任务是计算权重数字总和为质数的节点数。
例子：

Input: 
 

Output: 2 
Explanation: 
Node 1: digitSum(144) = 1 + 4 + 4 = 9 
Node 2: digitSum(1234) = 1 + 2 + 3 + 4 = 10 
Node 3: digitSum(21) = 2 + 1 = 3 
Node 4: digitSum(5) = 5 
Node 5: digitSum(77) = 7 + 7 = 14 
Only the sum of digits of the weights of nodes 3 and 4 are prime. 
 

方法：要解决上述问题，我们必须对树和每个节点执行DFS，并检查其权重的数字总和是否为质数。如果是，则增加计数。要检查数字总和是否为质数，我们将使用Eratosthenes筛。创建一个筛子，这将有助于我们确定度数是否在O(1)时间内为素数。
下面是上述方法的实现：

C++

// C++ program to Count Nodes
// which has Prime Digit
// Sum Weight in a Tree
 
#include 
using namespace std;
 
int MAX = 1000000;
int ans = 0;
 
vector graph[100];
vector weight(100);
 
// Function to create Sieve
// to check primes
void SieveOfEratosthenes(
    bool prime[], int p_size)
{
    // false here indicates
    // that it is not prime
    prime[0] = false;
    prime[1] = false;
 
    for (int p = 2; p * p <= p_size; p++) {
 
        // If prime[p] is not changed,
        // then it is a prime
        if (prime[p]) {
 
            // Update all multiples of p,
            // set them to non-prime
            for (int i = p * 2;
                 i <= p_size;
                 i += p)
                prime[i] = false;
        }
    }
}
 
// Function to return the
// sum of the digits of n
int digitSum(int n)
{
    int sum = 0;
    while (n) {
        sum += n % 10;
        n = n / 10;
    }
    return sum;
}
 
// Function to perform dfs
void dfs(int node,
         int parent,
         bool prime[])
{
    // If sum of the digits
    // of current node's weight
    // is prime then increment ans
    int sum = digitSum(weight[node]);
    if (prime[sum])
        ans += 1;
 
    for (int to : graph[node]) {
        if (to == parent)
            continue;
        dfs(to, node, prime);
    }
}
 
// Driver code
int main()
{
 
    // Weights of the node
    weight[1] = 144;
    weight[2] = 1234;
    weight[3] = 21;
    weight[4] = 5;
    weight[5] = 77;
 
    // Edges of the tree
    graph[1].push_back(2);
    graph[2].push_back(3);
    graph[2].push_back(4);
    graph[1].push_back(5);
 
    bool prime[MAX];
    memset(prime, true, sizeof(prime));
 
    SieveOfEratosthenes(prime, MAX);
 
    dfs(1, 1, prime);
 
    cout << ans;
 
    return 0;
}

Java

// Java program to Count Nodes
// which has Prime Digit
// Sum Weight in a Tree
import java.util.*;
class GFG{
 
static int MAX = 1000000;
static int ans = 0;
 
static Vector []graph =
              new Vector[100];
static int []weight = new int[100];
 
// Function to create Sieve
// to check primes
static void SieveOfEratosthenes(boolean prime[],
                                int p_size)
{
  // false here indicates
  // that it is not prime
  prime[0] = false;
  prime[1] = false;
 
  for (int p = 2; p * p <= p_size; p++)
  {
    // If prime[p] is not changed,
    // then it is a prime
    if (prime[p])
    {
      // Update all multiples of p,
      // set them to non-prime
      for (int i = p * 2;
               i < p_size; i += p)
        prime[i] = false;
    }
  }
}
 
// Function to return the
// sum of the digits of n
static int digitSum(int n)
{
  int sum = 0;
  while (n > 0)
  {
    sum += n % 10;
    n = n / 10;
  }
  return sum;
}
 
// Function to perform dfs
static void dfs(int node,
                int parent,
                boolean prime[])
{
  // If sum of the digits
  // of current node's weight
  // is prime then increment ans
  int sum = digitSum(weight[node]);
  if (prime[sum])
    ans += 1;
 
  for (int to : graph[node])
  {
    if (to == parent)
      continue;
    dfs(to, node, prime);
  }
}
 
// Driver code
public static void main(String[] args)
{
  // Weights of the node
  weight[1] = 144;
  weight[2] = 1234;
  weight[3] = 21;
  weight[4] = 5;
  weight[5] = 77;
  for (int i = 0; i < graph.length; i++)
    graph[i] = new Vector();
   
  // Edges of the tree
  graph[1].add(2);
  graph[2].add(3);
  graph[2].add(4);
  graph[1].add(5);
 
  boolean []prime = new boolean[MAX];
  Arrays.fill(prime, true);
  SieveOfEratosthenes(prime, MAX);
  dfs(1, 1, prime);
  System.out.print(ans);
}
}
 
// This code is contributed by Rajput-Ji

Python3

# Python program to Count Nodes
# which has Prime Digit
# Sum Weight in a Tree
from typing import List
MAX = 1000000
ans = 0
 
graph = [[] for _ in range(100)]
weight = [0 for _ in range(100)]
 
# Function to create Sieve
# to check primes
def SieveOfEratosthenes(prime: List[bool], p_size: int) -> None:
 
    # false here indicates
    # that it is not prime
    prime[0] = False
    prime[1] = False
 
    p = 2
    while p * p <= p_size:
       
        # If prime[p] is not changed,
        # then it is a prime
        if (prime[p]):
 
            # Update all multiples of p,
            # set them to non-prime
            for i in range(p * 2, p_size + 1, p):
                prime[i] = False
        p += 1
 
# Function to return the
# sum of the digits of n
def digitSum(n: int) -> int:
    sum = 0
    while (n):
        sum += n % 10
        n = n // 10
    return sum
 
# Function to perform dfs
def dfs(node: int, parent: int, prime: List[bool]) -> None:
    global ans
     
    # If sum of the digits
    # of current node's weight
    # is prime then increment ans
    sum = digitSum(weight[node])
    if (prime[sum]):
        ans += 1
    for to in graph[node]:
        if (to == parent):
            continue
        dfs(to, node, prime)
 
# Driver code
if __name__ == "__main__":
 
    # Weights of the node
    weight[1] = 144
    weight[2] = 1234
    weight[3] = 21
    weight[4] = 5
    weight[5] = 77
 
    # Edges of the tree
    graph[1].append(2)
    graph[2].append(3)
    graph[2].append(4)
    graph[1].append(5)
 
    prime = [True for _ in range(MAX + 1)]
    SieveOfEratosthenes(prime, MAX)
    dfs(1, 1, prime)
    print(ans)
 
# This code is contributed by sanjeev2552

C#

// C# program to Count Nodes
// which has Prime Digit
// Sum Weight in a Tree
using System;
using System.Collections.Generic;
class GFG{
 
static int MAX = 1000000;
static int ans = 0;
static List []graph =
       new List[100];
static int []weight = new int[100];
 
// Function to create Sieve
// to check primes
static void SieveOfEratosthenes(bool []prime,
                                int p_size)
{
  // false here indicates
  // that it is not prime
  prime[0] = false;
  prime[1] = false;
 
  for (int p = 2; p * p <= p_size; p++)
  {
    // If prime[p] is not changed,
    // then it is a prime
    if (prime[p])
    {
      // Update all multiples of p,
      // set them to non-prime
      for (int i = p * 2;
               i < p_size; i += p)
        prime[i] = false;
    }
  }
}
 
// Function to return the
// sum of the digits of n
static int digitSum(int n)
{
  int sum = 0;
  while (n > 0)
  {
    sum += n % 10;
    n = n / 10;
  }
  return sum;
}
 
// Function to perform dfs
static void dfs(int node,
                int parent,
                bool []prime)
{
  // If sum of the digits
  // of current node's weight
  // is prime then increment ans
  int sum = digitSum(weight[node]);
  if (prime[sum])
    ans += 1;
 
  foreach (int to in graph[node])
  {
    if (to == parent)
      continue;
    dfs(to, node, prime);
  }
}
 
// Driver code
public static void Main(String[] args)
{
  // Weights of the node
  weight[1] = 144;
  weight[2] = 1234;
  weight[3] = 21;
  weight[4] = 5;
  weight[5] = 77;
  for (int i = 0; i < graph.Length; i++)
    graph[i] = new List();
 
  // Edges of the tree
  graph[1].Add(2);
  graph[2].Add(3);
  graph[2].Add(4);
  graph[1].Add(5);
 
  bool []prime = new bool[MAX];
   
  for (int i = 0; i < prime.Length; i++)
    prime[i] = true;
   
  SieveOfEratosthenes(prime, MAX);
  dfs(1, 1, prime);
  Console.Write(ans);
}
}
 
// This code is contributed by Rajput-Ji

输出：

2

复杂度分析：

• 时间复杂度： O(N)。
  在DFS中，树的每个节点都处理一次，因此，如果树中总共有N个节点，则由于DFS导致的复杂度为O(N)。同样，为了处理每个节点，使用了SieveOfEratosthenes()函数，该函数的复杂度也为O(sqrt(N))，但由于此函数仅执行一次，因此不会影响整体时间复杂度。因此，时间复杂度为O(N)。
• 辅助空间： O(N)。
  素数数组使用了额外的空间，因此空间复杂度为O(N)。