Count Valid Paths in a Tree

Hard

#2564Time: 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).1 company

Patterns

Math Dynamic Programming Number Theory

Algorithms

Depth-First Search

Data structures

Tree

Companies

Sprinklr

Prompt

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] = [u_i, v_i] indicates that there is an edge between nodes u_i and v_i 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 <= 10⁵
edges.length == n - 1
edges[i].length == 2
1 <= u_i, v_i <= n
The input is generated such that edges represent a valid tree.

Approaches

2 approaches with complexity analysis and trade-offs.

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.

Algorithm

Create a boolean array isPrime of size n+1 and populate it using the Sieve of Eratosthenes.
Build an adjacency list for the tree from the edges array.
Initialize a global counter validPaths to 0.
Iterate through each node i from 1 to n.
For each i, start a DFS traversal dfs(u, parent, primeCount) from u=i, parent=-1, primeCount=0.
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).
After the outer loop finishes, return validPaths / 2 to correct for double counting.

Walkthrough

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.

1class Solution {2    private List<List<Integer>> adj;3    private boolean[] isPrime;4    private long count = 0;5    private int n;6 7    public long countPaths(int n, int[][] edges) {8        this.n = n;9        sieve(n);10        adj = new ArrayList<>();11        for (int i = 0; i <= n; i++) {12            adj.add(new ArrayList<>());13        }14        for (int[] edge : edges) {15            adj.get(edge[0]).add(edge[1]);16            adj.get(edge[1]).add(edge[0]);17        }18 19        for (int i = 1; i <= n; i++) {20            // We only need to start from a prime or a node adjacent to a prime21            // to find all paths, but for simplicity of brute force, we start from all.22            dfs(i, -1, 0);23        }24 25        return count / 2;26    }27 28    private void dfs(int u, int p, int primeCount) {29        if (isPrime[u]) {30            primeCount++;31        }32 33        if (primeCount > 1) {34            return;35        }36 37        // If start node is not u, and path is valid, count it.38        if (primeCount == 1) {39            count++;40        }41 42        for (int v : adj.get(u)) {43            if (v != p) {44                dfs(v, u, primeCount);45            }46        }47    }48 49    private void sieve(int n) {50        isPrime = new boolean[n + 1];51        if (n >= 2) {52            java.util.Arrays.fill(isPrime, true);53        }54        isPrime[0] = isPrime[1] = false;55        for (int i = 2; i * i <= n; i++) {56            if (isPrime[i]) {57                for (int j = i * i; j <= n; j += i) {58                    isPrime[j] = false;59                }60            }61        }62    }63}

Complexity

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

Trade-offs

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

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.

Algorithm

Generate primes up to n using Sieve of Eratosthenes.
Build an adjacency list for the tree.
Initialize a componentSize array and a visited array.
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.
The dfsComposite returns the size of the component. Store this size in the componentSize array for every node belonging to that component.
Initialize totalValidPaths = 0.
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.
Return totalValidPaths.

Walkthrough

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:

One endpoint is p, and the other is a composite node c in an adjacent composite component.
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.

1class Solution {2    private java.util.List<java.util.List<Integer>> adj;3    private boolean[] isPrime;4    private int[] componentSize;5    private boolean[] visited;6 7    public long countPaths(int n, int[][] edges) {8        sieve(n);9        adj = new java.util.ArrayList<>();10        for (int i = 0; i <= n; i++) {11            adj.add(new java.util.ArrayList<>());12        }13        for (int[] edge : edges) {14            adj.get(edge[0]).add(edge[1]);15            adj.get(edge[1]).add(edge[0]);16        }17 18        componentSize = new int[n + 1];19        visited = new boolean[n + 1];20        // Find sizes of connected components of composite numbers21        for (int i = 1; i <= n; i++) {22            if (!isPrime[i] && !visited[i]) {23                java.util.List<Integer> componentNodes = new java.util.ArrayList<>();24                int size = dfsComposite(i, componentNodes);25                for (int node : componentNodes) {26                    componentSize[node] = size;27                }28            }29        }30 31        long validPaths = 0;32        // Iterate through prime nodes to count paths33        for (int i = 1; i <= n; i++) {34            if (isPrime[i]) {35                long totalCompositeNeighbors = 0;36                java.util.List<Integer> neighborComponentSizes = new java.util.ArrayList<>();37                38                for (int neighbor : adj.get(i)) {39                    if (!isPrime[neighbor]) {40                        int size = componentSize[neighbor];41                        neighborComponentSizes.add(size);42                        totalCompositeNeighbors += size;43                    }44                }45                46                // Paths from prime 'i' to a composite node47                validPaths += totalCompositeNeighbors;48                49                // Paths between two different composite components connected via 'i'50                long sumOfProducts = 0;51                for (int size : neighborComponentSizes) {52                    sumOfProducts += (long)size * (totalCompositeNeighbors - size);53                }54                validPaths += sumOfProducts / 2;55            }56        }57 58        return validPaths;59    }60 61    private int dfsComposite(int u, java.util.List<Integer> componentNodes) {62        visited[u] = true;63        componentNodes.add(u);64        int count = 1;65        for (int v : adj.get(u)) {66            if (!isPrime[v] && !visited[v]) {67                count += dfsComposite(v, componentNodes);68            }69        }70        return count;71    }72 73    private void sieve(int n) {74        isPrime = new boolean[n + 1];75        if (n >= 2) {76            java.util.Arrays.fill(isPrime, true);77        }78        isPrime[0] = isPrime[1] = false;79        for (int i = 2; i * i <= n; i++) {80            if (isPrime[i]) {81                for (int j = i * i; j <= n; j += i) {82                    isPrime[j] = false;83                }84            }85        }86    }87}

Complexity

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.

Trade-offs

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.

Solutions

Solution

1class 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 ; } }

1const 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 ; } };

1class 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

Video walkthrough

Newsletter

One sharp idea, every week

System design and interview prep — short enough to finish.

No spam. Unsubscribe anytime.

Practice

Same difficulty — related problems to reinforce the pattern.

← Previous

Previous problem

Next problem