Count Valid Paths in a Tree

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

Examples

Example 1

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11



![](https://assets.leetcode.com/uploads/2023/08/27/example1.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

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13



![](https://assets.leetcode.com/uploads/2023/08/27/example2.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 <= 10^5
edges.length == n - 1
edges[i].length == 2
1 <= ui, vi <= n
The input is generated such that edges represent 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

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39


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

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47


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
}

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40


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

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42


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

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33


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

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48


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

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38


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.

Problem#

Examples#

Example 1#

Example 2#

Constraints#

Solution#

Method 1 – Prime Marking and DFS Subtree Counting#

Intuition#

Approach#

Code#

Complexity#