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

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.

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.

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