LeetCode, Hard++ (Acceptance 24%, Latest): 2867. Count Valid Paths in a Tree. DFS. O(n). Swift

Sergey Leschev. LeetCode Global TOP 200.
Sergey Leschev. LeetCode Global TOP 200.

Description

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 <= 10^5

edges.length == n - 1

edges[i].length == 2

1 <= u(i), v(i) <= n

The input is generated such that edges represent a valid tree.

Intuition

The intuition is to employ a depth-first search (DFS) approach through the tree.

Approach

During each step in the traversal, we perform the following key calculations:

  • Determine the path that ends at the current node.
  • Compute two different subtree paths that traverse the current node.
  • Maintain an array that keeps track of the cases where paths contain either 0 or 1 prime number.

This method allows us to efficiently count the valid paths in the tree while considering the presence of prime numbers.

Complexity

  • Time complexity: O(n)
  • Space complexity: O(n)

Code

class Solution { func countPaths(_ n: Int, _ edges: [[Int]]) -> Int { var ans = 0 var primes = Array(repeating: 1, count: n + 5) primes[1] = 0 for i in 2...Int(sqrt(Double(n + 5))) { if primes[i] == 0 { continue } for j in stride(from: i * i, to: primes.count, by: i) { primes[j] = 0 } } var adjList = Array(repeating: [Int](), count: n + 1) for edge in edges { adjList[edge[0]].append(edge[1]) adjList[edge[1]].append(edge[0]) } func dfs(_ curr: Int, _ parent: Int, _ adjList: [[Int]]) -> [Int] { var count = [0, 0] let isPrime = primes[curr] == 1 for child in adjList[curr] { if child == parent { continue } let next = dfs(child, curr, adjList) if isPrime { // Path ends at curr ans += next[0] // curr is conduit for path ans += count[1] * next[0] count[1] += next[0] } else { // Ends at curr ans += next[1] // curr is conduit for path ans += count[0] * next[1] ans += count[1] * next[0] count[0] += next[0] count[1] += next[1] } } count[isPrime ? 1 : 0] += 1 return count } dfs(1, -1, adjList) return ans } }

Sources: Github

Начать дискуссию