Skip to main content
Talentd

Count Valid Paths in a Tree - Solution & Explanation

HardMathDynamic ProgrammingTreeDepth-First Search29 min readAsked at: Gameskraft
Practice this problem

Problem Statement

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.

Approach Overview

Problem Overview: You are given a tree with n nodes labeled from 1..n. A path is valid if it contains exactly one prime-numbered node. The task is to count all such paths. Because the structure is a tree, every pair of nodes has exactly one simple path, which lets you reason about paths through prime nodes.

Approach 1: DFS Traversal with Path Prime Check (O(n log log n) time, O(n) space)

Start by precomputing all prime numbers up to n using the Sieve of Eratosthenes. The key observation: every valid path must include exactly one prime node, so treat each prime node as the “center” of potential paths. For each prime node, run DFS into its neighboring components but only traverse non-prime nodes. Each DFS returns the size of a connected component of non-prime nodes attached to that prime.

Suppose the component sizes are s1, s2, s3.... Any pair of nodes chosen from two different components forms a valid path that passes through the prime node. You also count single-ended paths from the prime to each non-prime node. Accumulate combinations using a running prefix sum to avoid quadratic merging. The algorithm performs a single traversal per component and uses adjacency lists for the tree. This approach relies heavily on Depth-First Search and basic Number Theory for prime detection.

Approach 2: Dynamic Programming on Tree (Tree DP) (O(n log log n) time, O(n) space)

This version computes path counts during a single DFS using dynamic programming on the tree. For every node, maintain two states: paths in its subtree that contain 0 primes and paths that contain 1 prime. When merging a child subtree into the parent, combine states carefully so the resulting path still contains exactly one prime.

If the current node is non-prime, zero-prime paths extend normally, and one-prime paths can be formed by attaching a subtree that already contains one prime. If the current node is prime, any extension immediately introduces the single allowed prime, so you only combine with zero-prime child paths. Each merge step contributes new valid paths to the global answer.

This technique avoids explicitly iterating over components for each prime node. Instead, it pushes information upward during DFS. The method is typical Dynamic Programming on a Tree, where each node summarizes information from its children before returning to its parent.

Recommended for interviews: The DFS component-count approach is the most common solution. It shows that you recognized the “exactly one prime” constraint and converted the problem into counting combinations of non-prime components around a prime node. The Tree DP approach is equally optimal but more complex to reason about under time pressure. Demonstrating the DFS solution first and discussing the DP variant shows strong algorithmic depth.

Approach 1: DFS Traversal with Path Prime Check

This approach uses Depth First Search (DFS) to traverse the tree and count valid paths. For each node, we check paths from the current node to all other nodes using DFS. We store whether a path contains a prime number. If exactly one node in the path is prime, it is a valid path. However, this method alone would be too slow for large trees, so optimizations around path sharing and careful counting are necessary.

The solution initializes a matrix to represent the adjacency list of the tree and an array to mark numbers as prime using the Sieve of Eratosthenes. We then use DFS to traverse all paths starting from node 1. It checks each path to see if it contains exactly one prime node.

Code

CC++JavaPythonC#JavaScript

C

C++

Java

Python

C#

JavaScript

Complexity

Time Complexity: O(n^2) due to the adjacency matrix and DFS traversal.
Space Complexity: O(n^2) for the adjacency matrix, which can be optimized by using adjacency list structure.

Try this approach in the editor →

Approach 2: Dynamic Programming on Tree (Tree DP)

This optimized approach uses Tree Dynamic Programming to calculate valid paths more efficiently. By employing tree-centric calculations, instead of evaluating every path individually, the algorithm calculates path properties at each node and aggregates results efficiently using dynamic programming techniques.

The C solution again uses a matrix for adjacency and checks subtree primes, optimizing path counting by focusing on direct connections. The program reduces repeated path checks by iterating through connections as opposed to paths.

Code

CC++JavaPythonC#JavaScript

C

C++

Java

Python

C#

JavaScript

Complexity

Time Complexity: O(n^2), though improved slightly in practice due to reduced path checks.
Space Complexity: O(n^2) similar to adjacency matrix range.

Try this approach in the editor →

Approach 3: Preprocessing + Union-Find + Enumeration

We can preprocess to get all the prime numbers in [1, n], where prime[i] indicates whether i is a prime number.

Next, we build a graph g based on the two-dimensional integer array, where g[i] represents all the neighbor nodes of node i. If both nodes of an edge are not prime numbers, we merge these two nodes into the same connected component.

Then, we enumerate all prime numbers i in the range of [1, n], considering all paths that include i.

Since i is already a prime number, if i is an endpoint of the path, we only need to accumulate the sizes of all connected components adjacent to node i. If i is a middle point on the path, we need to accumulate the product of the sizes of any two adjacent connected components.

The time complexity is O(n times \alpha(n)), and the space complexity is O(n). Here, n is the number of nodes, and \alpha is the inverse function of the Ackermann function.

Code

PythonJavaC++GoTypeScript

Python

Java

C++

Go

TypeScript

Try this approach in the editor →

Approach 4: Default Approach

Code

PythonJavaC++

Python

Java

C++

Try this approach in the editor →

Complexity Comparison

ApproachComplexity
DFS Traversal with Path Prime Check

Time Complexity: O(n^2) due to the adjacency matrix and DFS traversal.
Space Complexity: O(n^2) for the adjacency matrix, which can be optimized by using adjacency list structure.

Dynamic Programming on Tree (Tree DP)

Time Complexity: O(n^2), though improved slightly in practice due to reduced path checks.
Space Complexity: O(n^2) similar to adjacency matrix range.

Preprocessing + Union-Find + Enumeration—
Default Approach—

Detailed Complexity Analysis

ApproachTimeSpaceWhen to Use
DFS Traversal with Prime Component CountingO(n log log n)O(n)Best general solution. Clear reasoning: treat each prime node as a hub and count combinations of adjacent non-prime components.
Dynamic Programming on Tree (Tree DP)O(n log log n)O(n)Useful when solving multiple path-count variants in a tree or when you prefer single-pass DP aggregation.

Video Solution

2867. Count Valid Paths in a Tree | Weekly Leetcode 364 • codingMohan • 3,896 views views

Watch 7 more video solutions →

Frequently Asked Questions

Is Count Valid Paths in a Tree easy or hard?
Count Valid Paths in a Tree is classified as Hard. The challenge is recognizing that valid paths must contain exactly one prime node and converting the problem into counting combinations of non-prime components around each prime node using DFS or tree dynamic programming.
Count Valid Paths in a Tree Python/Java solution
Typical implementations build the adjacency list, compute primes with the Sieve of Eratosthenes, then run DFS from prime nodes to measure sizes of connected non-prime components. The same logic works in Python, Java, C++, and other languages with O(n log log n) time complexity.
How to solve Count Valid Paths in a Tree in O(n)?
Strict O(n) is difficult because prime detection typically uses a sieve costing O(n log log n). After preprocessing primes, the counting step is linear. DFS is used to measure sizes of non-prime components connected to each prime node, and combinations of those sizes produce the number of valid paths.
What is the best approach for Count Valid Paths in a Tree?
The most practical approach treats every prime node as the center of valid paths. Run DFS into adjacent non-prime components, compute their sizes, and count pair combinations across components. With prime preprocessing using a sieve, the total complexity becomes O(n log log n) time and O(n) space.
Is Count Valid Paths in a Tree asked at Google/Amazon/Meta?
Hard tree problems involving DFS, prime preprocessing, and path counting frequently appear in interviews at companies like Google, Amazon, and Meta. Variants that combine graph traversal with number theory or combinatorial counting are especially common in senior-level interview rounds.
What data structure is used in Count Valid Paths in a Tree?
The core data structure is an adjacency list to represent the tree efficiently. The algorithm also uses DFS recursion or a stack for traversal, arrays for prime flags from the sieve, and counters to track component sizes when combining paths.
What is the time complexity of Count Valid Paths in a Tree?
The optimal solutions run in O(n log log n) time due to the Sieve of Eratosthenes used to precompute primes. The DFS traversals themselves are linear, O(n), because each edge in the tree is processed a constant number of times. Space complexity is O(n) for adjacency lists and visited tracking.

Ready to solve this problem?

Practice Count Valid Paths in a Tree with our built-in code editor and test cases.

Practice on FleetCode
Previous

Count Paths That Can Form a Palindrome in a Tree