Scale Engineer

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
MathDynamic ProgrammingNumber 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] = [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.

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

  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.

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.

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

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

Solutions

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

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 →

Next problem