Count Valid Paths in a Tree
HardUpdated: Aug 2, 2025
Practice on:
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 thenumber 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 nodeaand ending with nodebsuch 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.
Examples
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^5edges.length == n - 1edges[i].length == 21 <= ui, vi <= n- The input is generated such that
edgesrepresent a valid tree.
Solution
Method 1 – Prime Marking and DFS Subtree Counting
Intuition
A valid path must contain exactly one prime-labeled node. For each prime node, count the number of nodes in its subtrees that are not prime, and use this to count valid paths passing through the prime. The answer is the sum over all primes of the number of pairs between the prime and each non-prime node in its connected components.
Approach
- Use the Sieve of Eratosthenes to mark all primes up to n.
- Build the tree as an adjacency list.
- For each prime node, for each connected component of non-prime nodes (after removing the prime), count the size of the component using DFS.
- For each such component, the number of valid paths is the size of the component (since each such node can pair with the prime node).
- Sum over all primes and all their non-prime components.
Code
C++
class Solution {
public:
long long countPaths(int n, vector<vector<int>>& edges) {
vector<bool> isPrime(n+1, 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;
}
}
vector<vector<int>> g(n+1);
for (auto& e : edges) {
g[e[0]].push_back(e[1]);
g[e[1]].push_back(e[0]);
}
long long ans = 0;
vector<bool> vis(n+1);
function<int(int)> dfs = [&](int u) {
vis[u] = true;
int cnt = 1;
for (int v : g[u]) {
if (!vis[v] && !isPrime[v]) cnt += dfs(v);
}
return cnt;
};
for (int i = 1; i <= n; ++i) {
if (!isPrime[i]) continue;
vis.assign(n+1, false);
vis[i] = true;
for (int v : g[i]) {
if (!isPrime[v] && !vis[v]) {
int cnt = dfs(v);
ans += cnt;
}
}
}
return ans;
}
};
Go
func countPaths(n int, edges [][]int) int64 {
isPrime := make([]bool, n+1)
for i := 2; i <= n; i++ {
isPrime[i] = true
}
for i := 2; i*i <= n; i++ {
if isPrime[i] {
for j := i * i; j <= n; j += i {
isPrime[j] = false
}
}
}
g := make([][]int, n+1)
for _, e := range edges {
g[e[0]] = append(g[e[0]], e[1])
g[e[1]] = append(g[e[1]], e[0])
}
var ans int64
vis := make([]bool, n+1)
var dfs func(int) int
dfs = func(u int) int {
vis[u] = true
cnt := 1
for _, v := range g[u] {
if !vis[v] && !isPrime[v] {
cnt += dfs(v)
}
}
return cnt
}
for i := 1; i <= n; i++ {
if !isPrime[i] {
continue
}
for j := range vis {
vis[j] = false
}
vis[i] = true
for _, v := range g[i] {
if !isPrime[v] && !vis[v] {
cnt := dfs(v)
ans += int64(cnt)
}
}
}
return ans
}
Java
class Solution {
public long countPaths(int n, int[][] edges) {
boolean[] isPrime = new boolean[n+1];
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;
}
}
List<Integer>[] g = new List[n+1];
for (int i = 0; i <= n; i++) g[i] = new ArrayList<>();
for (int[] e : edges) {
g[e[0]].add(e[1]);
g[e[1]].add(e[0]);
}
long ans = 0;
boolean[] vis = new boolean[n+1];
for (int i = 1; i <= n; ++i) {
if (!isPrime[i]) continue;
Arrays.fill(vis, false);
vis[i] = true;
for (int v : g[i]) {
if (!isPrime[v] && !vis[v]) {
int cnt = dfs(v, g, vis, isPrime);
ans += cnt;
}
}
}
return ans;
}
private int dfs(int u, List<Integer>[] g, boolean[] vis, boolean[] isPrime) {
vis[u] = true;
int cnt = 1;
for (int v : g[u]) {
if (!vis[v] && !isPrime[v]) cnt += dfs(v, g, vis, isPrime);
}
return cnt;
}
}
Kotlin
class Solution {
fun countPaths(n: Int, edges: Array<IntArray>): Long {
val isPrime = BooleanArray(n+1) { true }
isPrime[0] = false; isPrime[1] = false
for (i in 2..n) {
if (isPrime[i]) {
var j = i * i
while (j <= n) {
isPrime[j] = false
j += i
}
}
}
val g = Array(n+1) { mutableListOf<Int>() }
for (e in edges) {
g[e[0]].add(e[1])
g[e[1]].add(e[0])
}
var ans = 0L
val vis = BooleanArray(n+1)
fun dfs(u: Int): Int {
vis[u] = true
var cnt = 1
for (v in g[u]) {
if (!vis[v] && !isPrime[v]) cnt += dfs(v)
}
return cnt
}
for (i in 1..n) {
if (!isPrime[i]) continue
vis.fill(false)
vis[i] = true
for (v in g[i]) {
if (!isPrime[v] && !vis[v]) {
val cnt = dfs(v)
ans += cnt
}
}
}
return ans
}
}
Python
class Solution:
def countPaths(self, n: int, edges: list[list[int]]) -> int:
def sieve(n: int) -> list[bool]:
is_prime = [False, False] + [True] * (n-1)
for i in range(2, int(n**0.5)+1):
if is_prime[i]:
for j in range(i*i, n+1, i):
is_prime[j] = False
return is_prime
is_prime = sieve(n)
g = [[] for _ in range(n+1)]
for u, v in edges:
g[u].append(v)
g[v].append(u)
ans = 0
vis = [False] * (n+1)
def dfs(u: int) -> int:
vis[u] = True
cnt = 1
for v in g[u]:
if not vis[v] and not is_prime[v]:
cnt += dfs(v)
return cnt
for i in range(1, n+1):
if not is_prime[i]:
continue
vis = [False] * (n+1)
vis[i] = True
for v in g[i]:
if not is_prime[v] and not vis[v]:
cnt = dfs(v)
ans += cnt
return ans
Rust
impl Solution {
pub fn count_paths(n: i32, edges: Vec<Vec<i32>>) -> i64 {
let n = n as usize;
let mut is_prime = vec![true; n+1];
is_prime[0] = false;
is_prime[1] = false;
for i in 2..=n {
if is_prime[i] {
let mut j = i * i;
while j <= n {
is_prime[j] = false;
j += i;
}
}
}
let mut g = vec![vec![]; n+1];
for e in edges.iter() {
let u = e[0] as usize;
let v = e[1] as usize;
g[u].push(v);
g[v].push(u);
}
let mut ans = 0i64;
let mut vis = vec![false; n+1];
fn dfs(u: usize, g: &Vec<Vec<usize>>, vis: &mut Vec<bool>, is_prime: &Vec<bool>) -> i32 {
vis[u] = true;
let mut cnt = 1;
for &v in &g[u] {
if !vis[v] && !is_prime[v] {
cnt += dfs(v, g, vis, is_prime);
}
}
cnt
}
for i in 1..=n {
if !is_prime[i] { continue; }
vis.fill(false);
vis[i] = true;
for &v in &g[i] {
if !is_prime[v] && !vis[v] {
let cnt = dfs(v, &g, &mut vis, &is_prime);
ans += cnt as i64;
}
}
}
ans
}
}
TypeScript
class Solution {
countPaths(n: number, edges: number[][]): number {
const isPrime = Array(n+1).fill(true)
isPrime[0] = isPrime[1] = false
for (let i = 2; i*i <= n; ++i) {
if (isPrime[i]) {
for (let j = i*i; j <= n; j += i) isPrime[j] = false
}
}
const g: number[][] = Array.from({length: n+1}, () => [])
for (const [u, v] of edges) {
g[u].push(v)
g[v].push(u)
}
let ans = 0
const vis = Array(n+1).fill(false)
function dfs(u: number): number {
vis[u] = true
let cnt = 1
for (const v of g[u]) {
if (!vis[v] && !isPrime[v]) cnt += dfs(v)
}
return cnt
}
for (let i = 1; i <= n; ++i) {
if (!isPrime[i]) continue
vis.fill(false)
vis[i] = true
for (const v of g[i]) {
if (!isPrime[v] && !vis[v]) {
const cnt = dfs(v)
ans += cnt
}
}
}
return ans
}
}
Complexity
- ⏰ Time complexity:
O(n)for sieve and tree traversal, each node and edge is visited a constant number of times. - 🧺 Space complexity:
O(n)for the tree, sieve, and visited arrays.