2867. Count Valid Paths in a Tree LeetCode Solution

In this guide, you will get 2867. Count Valid Paths in a Tree LeetCode Solution with the best time and space complexity. The solution to Count Valid Paths in a Tree problem is provided in various programming languages like C++, Java, and Python. This will be helpful for you if you are preparing for placements, hackathons, interviews, or practice purposes. The solutions provided here are very easy to follow and include detailed explanations.

Table of Contents

  1. Problem Statement
  2. Complexity Analysis
  3. Count Valid Paths in a Tree solution in C++
  4. Count Valid Paths in a Tree solution in Java
  5. Count Valid Paths in a Tree solution in Python
  6. Additional Resources
2867. Count Valid Paths in a Tree LeetCode Solution image

Problem Statement of Count Valid Paths in a Tree

There is an undirected tree with n nodes labeled from 1 to n. You are given the integer n and a 2D integer array edges of length n – 1, where edges[i] = [ui, vi] indicates that there is an edge between nodes ui and vi in the tree.
Return the number of valid paths in the tree.
A path (a, b) is valid if there exists exactly one prime number among the node labels in the path from a to b.
Note that:

The path (a, b) is a sequence of distinct nodes starting with node a and ending with node b such that every two adjacent nodes in the sequence share an edge in the tree.
Path (a, b) and path (b, a) are considered the same and counted only once.

Example 1:

Input: n = 5, edges = [[1,2],[1,3],[2,4],[2,5]]
Output: 4
Explanation: The pairs with exactly one prime number on the path between them are:
– (1, 2) since the path from 1 to 2 contains prime number 2.
– (1, 3) since the path from 1 to 3 contains prime number 3.
– (1, 4) since the path from 1 to 4 contains prime number 2.
– (2, 4) since the path from 2 to 4 contains prime number 2.
It can be shown that there are only 4 valid paths.

Example 2:

Input: n = 6, edges = [[1,2],[1,3],[2,4],[3,5],[3,6]]
Output: 6
Explanation: The pairs with exactly one prime number on the path between them are:
– (1, 2) since the path from 1 to 2 contains prime number 2.
– (1, 3) since the path from 1 to 3 contains prime number 3.
– (1, 4) since the path from 1 to 4 contains prime number 2.
– (1, 6) since the path from 1 to 6 contains prime number 3.
– (2, 4) since the path from 2 to 4 contains prime number 2.
– (3, 6) since the path from 3 to 6 contains prime number 3.
It can be shown that there are only 6 valid paths.

Constraints:

1 <= n <= 105
edges.length == n – 1
edges[i].length == 2
1 <= ui, vi <= n
The input is generated such that edges represent a valid tree.

Complexity Analysis

  • Time Complexity: O(n)
  • Space Complexity: O(n)

2867. Count Valid Paths in a Tree LeetCode Solution in C++

class Solution {
 public:
  long long countPaths(int n, vector<vector<int>>& edges) {
    long ans = 0;
    const vector<bool> isPrime = sieveEratosthenes(n + 1);
    vector<vector<int>> graph(n + 1);

    for (const vector<int>& edge : edges) {
      const int u = edge[0];
      const int v = edge[1];
      graph[u].push_back(v);
      graph[v].push_back(u);
    }

    dfs(graph, 1, /*prev=*/-1, isPrime, ans);
    return ans;
  }

 private:
  pair<long, long> dfs(const vector<vector<int>>& graph, int u, int prev,
                       const vector<bool>& isPrime, long& ans) {
    long countZeroPrimePath = !isPrime[u];
    long countOnePrimePath = isPrime[u];

    for (const int v : graph[u]) {
      if (v == prev)
        continue;
      const auto& [countZeroPrimeChildPath, countOnePrimeChildPath] =
          dfs(graph, v, u, isPrime, ans);
      ans += countZeroPrimePath * countOnePrimeChildPath +
             countOnePrimePath * countZeroPrimeChildPath;
      if (isPrime[u]) {
        countOnePrimePath += countZeroPrimeChildPath;
      } else {
        countZeroPrimePath += countZeroPrimeChildPath;
        countOnePrimePath += countOnePrimeChildPath;
      }
    }

    return {countZeroPrimePath, countOnePrimePath};
  }

  vector<bool> sieveEratosthenes(int n) {
    vector<bool> isPrime(n, true);
    isPrime[0] = false;
    isPrime[1] = false;
    for (int i = 2; i * i < n; ++i)
      if (isPrime[i])
        for (int j = i * i; j < n; j += i)
          isPrime[j] = false;
    return isPrime;
  }
};
/* code provided by PROGIEZ */

2867. Count Valid Paths in a Tree LeetCode Solution in Java

class Solution {
  public long countPaths(int n, int[][] edges) {
    final boolean[] isPrime = sieveEratosthenes(n + 1);
    List<Integer>[] graph = new List[n + 1];

    for (int i = 1; i <= n; ++i)
      graph[i] = new ArrayList<>();

    for (int[] edge : edges) {
      final int u = edge[0];
      final int v = edge[1];
      graph[u].add(v);
      graph[v].add(u);
    }

    dfs(graph, 1, /*prev=*/-1, isPrime);
    return ans;
  }

  private long ans = 0;

  private Pair<Long, Long> dfs(List<Integer>[] graph, int u, int prev, boolean[] isPrime) {
    long countZeroPrimePath = isPrime[u] ? 0 : 1;
    long countOnePrimePath = isPrime[u] ? 1 : 0;

    for (final int v : graph[u]) {
      if (v == prev)
        continue;
      Pair<Long, Long> pair = dfs(graph, v, u, isPrime);
      final long countZeroPrimeChildPath = pair.getKey();
      final long countOnePrimeChildPath = pair.getValue();
      ans +=
          countZeroPrimePath * countOnePrimeChildPath + countOnePrimePath * countZeroPrimeChildPath;
      if (isPrime[u]) {
        countOnePrimePath += countZeroPrimeChildPath;
      } else {
        countZeroPrimePath += countZeroPrimeChildPath;
        countOnePrimePath += countOnePrimeChildPath;
      }
    }

    return new Pair<>(countZeroPrimePath, countOnePrimePath);
  }

  private boolean[] sieveEratosthenes(int n) {
    boolean[] isPrime = new boolean[n];
    Arrays.fill(isPrime, true);
    isPrime[0] = false;
    isPrime[1] = false;
    for (int i = 2; i * i < n; ++i)
      if (isPrime[i])
        for (int j = i * i; j < n; j += i)
          isPrime[j] = false;
    return isPrime;
  }
}
// code provided by PROGIEZ

2867. Count Valid Paths in a Tree LeetCode Solution in Python

class Solution:
  def countPaths(self, n: int, edges: list[list[int]]) -> int:
    ans = 0
    isPrime = self._sieveEratosthenes(n + 1)
    graph = [[] for _ in range(n + 1)]

    for u, v in edges:
      graph[u].append(v)
      graph[v].append(u)

    def dfs(u: int, prev: int) -> tuple[int, int]:
      nonlocal ans
      countZeroPrimePath = int(not isPrime[u])
      countOnePrimePath = int(isPrime[u])

      for v in graph[u]:
        if v == prev:
          continue
        countZeroPrimeChildPath, countOnePrimeChildPath = dfs(v, u)
        ans += (countZeroPrimePath * countOnePrimeChildPath +
                countOnePrimePath * countZeroPrimeChildPath)
        if isPrime[u]:
          countOnePrimePath += countZeroPrimeChildPath
        else:
          countZeroPrimePath += countZeroPrimeChildPath
          countOnePrimePath += countOnePrimeChildPath

      return countZeroPrimePath, countOnePrimePath

    dfs(1, -1)
    return ans

  def _sieveEratosthenes(self, n: int) -> list[bool]:
    isPrime = [True] * n
    isPrime[0] = False
    isPrime[1] = False
    for i in range(2, int(n**0.5) + 1):
      if isPrime[i]:
        for j in range(i * i, n, i):
          isPrime[j] = False
    return isPrime
# code by PROGIEZ

Additional Resources

Happy Coding! Keep following PROGIEZ for more updates and solutions.