#2867

Hard

auto_awesomeBinary-tree traversal and state tracking

LeetCode Problem Workspace

Count Valid Paths in a Tree

Count Valid Paths in a Tree involves finding paths with exactly one prime number in a tree of n nodes.

math dynamic programming tree depth first search number theory

Problem Statement

You are given a tree with n nodes labeled from 1 to n. The tree is represented by an array of edges, where each edge connects two nodes. Your task is to find and return the number of valid paths in the tree. A path is valid if it contains exactly one prime number among the node labels.

The input consists of the number of nodes n and a 2D integer array edges, where edges[i] = [ui, vi] denotes an edge between nodes ui and vi. You must determine the number of valid paths that can be formed between any two nodes in the tree, following the specific prime number constraint.

Examples

Example 1

Input: n = 5, edges = [[1,2],[1,3],[2,4],[2,5]]

Output: 4

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

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.

Solution Approach

Prime Number Calculation

Use the Sieve of Eratosthenes to find all prime numbers up to n. This preprocessing step ensures that you can quickly check whether a node label is prime during traversal.

Tree Traversal with DFS

Perform a Depth-First Search (DFS) traversal of the tree. As you visit each node, track the count of prime numbers encountered along the path from the root to that node. This will allow efficient validation of paths between any two nodes.

State Tracking with Dynamic Programming

Use dynamic programming to store the prime count and path information at each node. This approach reduces redundant calculations by ensuring that previously computed results are reused during the DFS.

Complexity Analysis

Metric	Value
Time	Depends on the final approach
Space	Depends on the final approach

Time and space complexity depend on the chosen approach. The sieve of Eratosthenes runs in O(n log log n), while the DFS traversal and dynamic programming state updates are O(n), making the overall time complexity approximately O(n log log n). Space complexity is O(n) due to the storage requirements for the sieve and dynamic programming table.

What Interviewers Usually Probe

Check if the candidate can optimize tree traversal using DFS and dynamic programming.
Look for the ability to handle large inputs efficiently, using prime number computation and state tracking.
Evaluate how the candidate deals with the prime number constraint, especially in edge cases where there are few or no primes.

Common Pitfalls or Variants

Common pitfalls

Forgetting to mark non-prime numbers correctly in the sieve of Eratosthenes, leading to incorrect path evaluations.
Inefficient tree traversal that does not minimize redundant work, especially in larger trees.
Incorrectly counting paths that contain more than one prime number, violating the problem's constraints.

Follow-up variants

Consider trees with only one node, where no paths exist.
Modify the problem to count paths with exactly two primes instead of one.
Explore how the solution scales when the tree is unbalanced or contains many leaves.

FAQ

What is the primary pattern for solving 'Count Valid Paths in a Tree'?

The primary pattern is binary-tree traversal with state tracking, using dynamic programming to efficiently count valid paths based on prime number constraints.

How do you efficiently find prime numbers for this problem?

Use the Sieve of Eratosthenes to compute all prime numbers up to n before starting the tree traversal. This allows quick checks during the DFS.

What is the time complexity of solving this problem?

The time complexity is O(n log log n) due to the sieve of Eratosthenes, with DFS traversal and dynamic programming updates each running in O(n).

Can this problem be solved with a BFS approach?

Yes, while DFS is commonly used, BFS can also be employed for tree traversal. However, DFS is generally more suited for depth-based path tracking.

How does GhostInterview help with this problem?

GhostInterview guides you through the key steps, helping you optimize tree traversal and understand prime number constraints, while providing targeted feedback on common pitfalls.

terminal

Solution

Solution 1: 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.

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

49

50

51

52

53

class 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

Solution 2

#### Python3

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

49

50

51

52

53

class 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

Continue Practicing

#3559 number of ways to assign edge weights ii #2920 maximum points after collecting coins from all nodes #2925 maximum score after applying operations on a tree