# Count Valid Paths in a Tree
**Difficulty:** HARD
[External](https://leetcode.com/problems/count-valid-paths-in-a-tree)
Canonical: https://scaleengineer.com/dsa/problems/count-valid-paths-in-a-tree
**Patterns:** [Math](https://scaleengineer.com/dsa/patterns/math), [Dynamic Programming](https://scaleengineer.com/dsa/patterns/dynamic-programming), [Number Theory](https://scaleengineer.com/dsa/patterns/number-theory)
**Algorithms:** [Depth-First Search](https://scaleengineer.com/algorithms/depth-first-search)
**Data structures:** Tree
**Companies:** [Sprinklr](https://scaleengineer.com/companies/sprinklr)
---
## Problem
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:**

![](https://assets.glich.co/dsa/count-valid-paths-in-a-tree/image0.png) 

**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:**

![](https://assets.glich.co/dsa/count-valid-paths-in-a-tree/image1.png) 

**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.

# Approaches
## Brute-Force Traversal from Each Node
This approach iterates through every possible starting node in the tree. For each starting node, it performs a traversal (like Depth-First Search or Breadth-First Search) to find all possible paths originating from it. During the traversal, it keeps track of the number of prime nodes encountered on the current path. If a path to a destination node contains exactly one prime number, it's counted as a valid path.
**Time:** O(N^2). The outer loop runs N times. Inside, the DFS can traverse up to N nodes. This leads to N * O(N) = O(N^2) complexity. The initial Sieve takes O(N log log N), but it's dominated by the main loop. · **Space:** O(N). We need space for the adjacency list (O(N)), the `isPrime` array (O(N)), and the recursion stack for DFS (O(N) in the worst case for a skewed tree).
**Pros:** Relatively simple to understand and implement compared to more optimized solutions.; Correctly solves the problem for small `n`.
**Cons:** Highly inefficient due to redundant computations. The traversal from each node re-explores large parts of the tree.; Will result in a "Time Limit Exceeded" error for the given constraints (`n` up to 10^5).
### Explanation
First, we need a way to quickly check if a number is prime. We can pre-compute primality for all numbers from 1 to `n` using the Sieve of Eratosthenes and store the results in a boolean array.
Next, we build an adjacency list representation of the tree from the given `edges`.
The main part of the algorithm involves a nested loop. The outer loop iterates through each node `i` from 1 to `n`, considering it as a potential starting point of a path.
From each starting node `i`, we initiate a Depth-First Search (DFS). The DFS function, say `dfs(u, p, primeCount)`, will explore the tree, where `u` is the current node, `p` is its parent (to avoid going backward), and `primeCount` is the count of prime numbers on the path from the starting node `i` to `u`.
Inside the DFS, we first update `primeCount` based on whether the current node `u` is prime.
If the `primeCount` for the path from `i` to `u` is exactly 1, we've found a valid path `(i, u)`. We've found a valid path `(i, u)`. We increment our total count.
The DFS continues recursively to all neighbors of `u` (except its parent), but only if the `primeCount` does not exceed 1. If it's already 2 or more, any further extension of the path will also have at least 2 primes, so we can prune this search branch.
Since this process counts each path `(a, b)` twice (once when starting from `a` and once from `b`), the final result is the total count divided by 2.
```java
class Solution {
    private List<List<Integer>> adj;
    private boolean[] isPrime;
    private long count = 0;
    private int n;

    public long countPaths(int n, int[][] edges) {
        this.n = n;
        sieve(n);
        adj = new ArrayList<>();
        for (int i = 0; i <= n; i++) {
            adj.add(new ArrayList<>());
        }
        for (int[] edge : edges) {
            adj.get(edge[0]).add(edge[1]);
            adj.get(edge[1]).add(edge[0]);
        }

        for (int i = 1; i <= n; i++) {
            // We only need to start from a prime or a node adjacent to a prime
            // to find all paths, but for simplicity of brute force, we start from all.
            dfs(i, -1, 0);
        }

        return count / 2;
    }

    private void dfs(int u, int p, int primeCount) {
        if (isPrime[u]) {
            primeCount++;
        }

        if (primeCount > 1) {
            return;
        }

        // If start node is not u, and path is valid, count it.
        if (primeCount == 1) {
            count++;
        }

        for (int v : adj.get(u)) {
            if (v != p) {
                dfs(v, u, primeCount);
            }
        }
    }

    private void sieve(int n) {
        isPrime = new boolean[n + 1];
        if (n >= 2) {
            java.util.Arrays.fill(isPrime, true);
        }
        isPrime[0] = 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;
                }
            }
        }
    }
}
```
### Algorithm
1. Create a boolean array `isPrime` of size `n+1` and populate it using the Sieve of Eratosthenes.
2. Build an adjacency list for the tree from the `edges` array.
3. Initialize a global counter `validPaths` to 0.
4. Iterate through each node `i` from 1 to `n`.
5. For each `i`, start a DFS traversal `dfs(u, parent, primeCount)` from `u=i`, `parent=-1`, `primeCount=0`.
6. In the `dfs(u, p, pc)` function:
    a. Update the prime count: `new_pc = pc + (isPrime[u] ? 1 : 0)`.
    b. If `new_pc > 1`, stop exploring this path and return.
    c. If `new_pc == 1`, it means the path from the start node `i` to `u` is valid. Increment `validPaths`.
    d. For each neighbor `v` of `u` (where `v != p`), recursively call `dfs(v, u, new_pc)`.
7. After the outer loop finishes, return `validPaths / 2` to correct for double counting.

## Counting Paths via Composite Components
This efficient approach is based on a key observation: any valid path must contain exactly one prime node. This means prime nodes act as "bridges" between components formed exclusively by composite nodes. By removing all prime nodes, the tree splits into a forest of smaller trees (components) where every node is composite. We can then count valid paths by considering how these composite components are connected through prime nodes.
**Time:** O(N log log N). The Sieve of Eratosthenes takes O(N log log N). Building the adjacency list is O(N). Finding composite components is O(N) as it's a single traversal over the composite subgraph. The final loop iterates through primes and their neighbors, and the sum of degrees is 2*(N-1), so this step is also O(N). The bottleneck is the Sieve. · **Space:** O(N). For the `isPrime` array, adjacency list, `visited` array, `componentSize` array, and recursion stack.
**Pros:** Highly efficient and optimal.; Avoids redundant calculations by processing each part of the graph (composite components, prime node connections) only once.; Passes for large constraints on `n`.
**Cons:** More complex to conceptualize and implement than the brute-force approach.; Requires careful handling of counting to avoid overcounting or undercounting.
### Explanation
The algorithm begins by pre-computing prime numbers up to `n` using a Sieve.
We then build an adjacency list for the tree.
The core idea is to first identify the connected components formed by only composite nodes. We can do this with a single pass of DFS or BFS. We traverse the graph, but only move between composite nodes. For each composite component we find, we calculate its size (number of nodes) and store this information. A `componentSize` array can map each composite node to the size of its component.
After identifying the composite components, we iterate through each prime node `p` in the tree. For each prime `p`, we count the number of valid paths that have `p` as their *sole* prime member.
A path involving `p` is valid if it's of the form `(a, b)` where the path from `a` to `b` contains only `p` as a prime. This can happen in two ways:
1. One endpoint is `p`, and the other is a composite node `c` in an adjacent composite component.
2. Both endpoints `a` and `b` are composite nodes, but they belong to *different* composite components that are both adjacent to `p`.
To count these, for a prime `p`, we look at its neighbors. For each neighbor `v` that is composite, we find the size of the composite component it belongs to (which we pre-calculated). Let these sizes be `s_1, s_2, ..., s_k`.
The total number of composite nodes reachable from `p` without crossing another prime is `total_neighbors = s_1 + s_2 + ... + s_k`.
The number of valid paths of type 1 (from `p` to a composite) is simply `total_neighbors`.
The number of valid paths of type 2 (between two different composite components) is the sum of `s_i * s_j` for all pairs of distinct components `i` and `j`. A simpler way to calculate this is `( (s_1* (total_neighbors - s_1)) + (s_2 * (total_neighbors - s_2)) + ... ) / 2`.
The total paths for prime `p` is the sum of these two counts. We sum this value over all prime nodes to get the final answer. Since each valid path has a unique prime, this process counts every valid path exactly once.
```java
class Solution {
    private java.util.List<java.util.List<Integer>> adj;
    private boolean[] isPrime;
    private int[] componentSize;
    private boolean[] visited;

    public long countPaths(int n, int[][] edges) {
        sieve(n);
        adj = new java.util.ArrayList<>();
        for (int i = 0; i <= n; i++) {
            adj.add(new java.util.ArrayList<>());
        }
        for (int[] edge : edges) {
            adj.get(edge[0]).add(edge[1]);
            adj.get(edge[1]).add(edge[0]);
        }

        componentSize = new int[n + 1];
        visited = new boolean[n + 1];
        // Find sizes of connected components of composite numbers
        for (int i = 1; i <= n; i++) {
            if (!isPrime[i] && !visited[i]) {
                java.util.List<Integer> componentNodes = new java.util.ArrayList<>();
                int size = dfsComposite(i, componentNodes);
                for (int node : componentNodes) {
                    componentSize[node] = size;
                }
            }
        }

        long validPaths = 0;
        // Iterate through prime nodes to count paths
        for (int i = 1; i <= n; i++) {
            if (isPrime[i]) {
                long totalCompositeNeighbors = 0;
                java.util.List<Integer> neighborComponentSizes = new java.util.ArrayList<>();
                
                for (int neighbor : adj.get(i)) {
                    if (!isPrime[neighbor]) {
                        int size = componentSize[neighbor];
                        neighborComponentSizes.add(size);
                        totalCompositeNeighbors += size;
                    }
                }
                
                // Paths from prime 'i' to a composite node
                validPaths += totalCompositeNeighbors;
                
                // Paths between two different composite components connected via 'i'
                long sumOfProducts = 0;
                for (int size : neighborComponentSizes) {
                    sumOfProducts += (long)size * (totalCompositeNeighbors - size);
                }
                validPaths += sumOfProducts / 2;
            }
        }

        return validPaths;
    }

    private int dfsComposite(int u, java.util.List<Integer> componentNodes) {
        visited[u] = true;
        componentNodes.add(u);
        int count = 1;
        for (int v : adj.get(u)) {
            if (!isPrime[v] && !visited[v]) {
                count += dfsComposite(v, componentNodes);
            }
        }
        return count;
    }

    private void sieve(int n) {
        isPrime = new boolean[n + 1];
        if (n >= 2) {
            java.util.Arrays.fill(isPrime, true);
        }
        isPrime[0] = 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;
                }
            }
        }
    }
}
```
### Algorithm
1. Generate primes up to `n` using Sieve of Eratosthenes.
2. Build an adjacency list for the tree.
3. Initialize a `componentSize` array and a `visited` array.
4. Iterate through all nodes from 1 to `n`. If a node `i` is composite and not visited, start a DFS (`dfsComposite`) to find all nodes in its connected component of composite numbers.
5. The `dfsComposite` returns the size of the component. Store this size in the `componentSize` array for every node belonging to that component.
6. Initialize `totalValidPaths = 0`.
7. Iterate through each node `p` from 1 to `n`. If `p` is prime:
    a. Initialize `total_composite_neighbors = 0` and a list `neighbor_component_sizes`.
    b. For each neighbor `v` of `p`:
        i. If `v` is composite, get its component size `s = componentSize[v]`.
        ii. Add `s` to `neighbor_component_sizes`.
        iii. Add `s` to `total_composite_neighbors`.
    c. Add `total_composite_neighbors` to `totalValidPaths` (paths from `p` to a composite node).
    d. Calculate paths between different components: iterate through `neighbor_component_sizes`. For each size `s`, add `s * (total_composite_neighbors - s)` to a temporary sum.
    e. Add `(temporary sum) / 2` to `totalValidPaths`.
8. Return `totalValidPaths`.

# Solutions
### Java

```java
class PrimeTable { private final boolean [] prime ; public PrimeTable ( int n ) { prime = new boolean [ n + 1 ]; Arrays . fill ( prime , true ); prime [ 0 ] = false ; prime [ 1 ] = false ; for ( int i = 2 ; i <= n ; ++ i ) { if ( prime [ i ]) { for ( int j = i + i ; j <= n ; j += i ) { prime [ j ] = false ; } } } } public boolean isPrime ( int x ) { return prime [ x ]; } } class UnionFind { private final int [] p ; private final int [] size ; public UnionFind ( int n ) { p = new int [ n ]; size = new int [ n ]; for ( int i = 0 ; i < n ; ++ i ) { p [ i ] = i ; size [ i ] = 1 ; } } public int find ( int x ) { if ( p [ x ] != x ) { p [ x ] = find ( p [ x ]); } return p [ x ]; } public boolean union ( int a , int b ) { int pa = find ( a ), pb = find ( b ); if ( pa == pb ) { return false ; } if ( size [ pa ] > size [ pb ]) { p [ pb ] = pa ; size [ pa ] += size [ pb ]; } else { p [ pa ] = pb ; size [ pb ] += size [ pa ]; } return true ; } public int size ( int x ) { return size [ find ( x )]; } } class Solution { private static final PrimeTable PT = new PrimeTable ( 100010 ); public long countPaths ( int n , int [][] edges ) { List < Integer >[] g = new List [ n + 1 ]; Arrays . setAll ( g , i -> new ArrayList <>()); UnionFind uf = new UnionFind ( n + 1 ); for ( int [] e : edges ) { int u = e [ 0 ], v = e [ 1 ]; g [ u ]. add ( v ); g [ v ]. add ( u ); if (! PT . isPrime ( u ) && ! PT . isPrime ( v )) { uf . union ( u , v ); } } long ans = 0 ; for ( int i = 1 ; i <= n ; ++ i ) { if ( PT . isPrime ( i )) { long t = 0 ; for ( int j : g [ i ]) { if (! PT . isPrime ( j )) { long cnt = uf . size ( j ); ans += cnt ; ans += cnt * t ; t += cnt ; } } } } return ans ; } }
```

### CPP

```cpp
const int mx = 1e5 + 10 ; bool prime [ mx + 1 ]; int init = []() { for ( int i = 2 ; i <= mx ; ++ i ) prime [ i ] = true ; for ( int i = 2 ; i <= mx ; ++ i ) { if ( prime [ i ]) { for ( int j = i + i ; j <= mx ; j += i ) { prime [ j ] = false ; } } } return 0 ; }(); class UnionFind { public: UnionFind ( int n ) { p = vector < int > ( n ); size = vector < int > ( n , 1 ); iota ( p . begin (), p . end (), 0 ); } bool unite ( int a , int b ) { int pa = find ( a ), pb = find ( b ); if ( pa == pb ) { return false ; } if ( size [ pa ] > size [ pb ]) { p [ pb ] = pa ; size [ pa ] += size [ pb ]; } else { p [ pa ] = pb ; size [ pb ] += size [ pa ]; } return true ; } int find ( int x ) { if ( p [ x ] != x ) { p [ x ] = find ( p [ x ]); } return p [ x ]; } int getSize ( int x ) { return size [ find ( x )]; } private: vector < int > p , size ; }; class Solution { public: long long countPaths ( int n , vector < vector < int >>& edges ) { vector < int > g [ n + 1 ]; UnionFind uf ( n + 1 ); for ( auto & e : edges ) { int u = e [ 0 ], v = e [ 1 ]; g [ u ]. push_back ( v ); g [ v ]. push_back ( u ); if ( ! prime [ u ] && ! prime [ v ]) { uf . unite ( u , v ); } } long long ans = 0 ; for ( int i = 1 ; i <= n ; ++ i ) { if ( prime [ i ]) { long long t = 0 ; for ( int j : g [ i ]) { if ( ! prime [ j ]) { long long cnt = uf . getSize ( j ); ans += cnt ; ans += cnt * t ; t += cnt ; } } } } return ans ; } };
```

### Python

```python
class UnionFind : def __init__ ( self , n ): self . p = list ( range ( n )) self . size = [ 1 ] * n def find ( self , x ): if self . p [ x ] != x : self . p [ x ] = self . find ( self . p [ x ]) return self . p [ x ] def union ( self , a , b ): pa , pb = self . find ( a ), self . find ( b ) if pa == pb : return False if self . size [ pa ] > self . size [ pb ]: self . p [ pb ] = pa self . size [ pa ] += self . size [ pb ] else : self . p [ pa ] = pb self . size [ pb ] += self . size [ pa ] return True mx = 10 ** 5 + 10 prime = [ True ] * ( mx + 1 ) prime [ 0 ] = prime [ 1 ] = False for i in range ( 2 , mx + 1 ): if prime [ i ]: for j in range ( i * i , mx + 1 , i ): prime [ j ] = False class Solution : def countPaths ( self , n : int , edges : List [ List [ int ]]) -> int : g = [[] for _ in range ( n + 1 )] uf = UnionFind ( n + 1 ) for u , v in edges : g [ u ]. append ( v ) g [ v ]. append ( u ) if prime [ u ] + prime [ v ] == 0 : uf . union ( u , v ) ans = 0 for i in range ( 1 , n + 1 ): if prime [ i ]: t = 0 for j in g [ i ]: if not prime [ j ]: cnt = uf . size [ uf . find ( j )] ans += cnt ans += t * cnt t += cnt return ans
```