<p>There is an undirected tree with <code>n</code> nodes labeled from <code>1</code> to <code>n</code>. You are given the integer <code>n</code> and a 2D integer array <code>edges</code> of length <code>n - 1</code>, where <code>edges[i] = [u<sub>i</sub>, v<sub>i</sub>]</code> indicates that there is an edge between nodes <code>u<sub>i</sub></code> and <code>v<sub>i</sub></code> in the tree.</p>

<p>Return <em>the <strong>number of valid paths</strong> in the tree</em>.</p>

<p>A path <code>(a, b)</code> is <strong>valid</strong> if there exists <strong>exactly one</strong> prime number among the node labels in the path from <code>a</code> to <code>b</code>.</p>

<p><strong>Note</strong> that:</p>

<ul>
	<li>The path <code>(a, b)</code> is a sequence of <strong>distinct</strong> nodes starting with node <code>a</code> and ending with node <code>b</code> such that every two adjacent nodes in the sequence share an edge in the tree.</li>
	<li>Path <code>(a, b)</code> and path <code>(b, a)</code> are considered the <strong>same</strong> and counted only <strong>once</strong>.</li>
</ul>

<p>&nbsp;</p>
<p><strong class="example">Example 1:</strong></p>
<img alt="" src="https://assets.leetcode.com/uploads/2023/08/27/example1.png" style="width: 440px; height: 357px;" />
<pre>
<strong>Input:</strong> n = 5, edges = [[1,2],[1,3],[2,4],[2,5]]
<strong>Output:</strong> 4
<strong>Explanation:</strong> 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.
</pre>

<p><strong class="example">Example 2:</strong></p>
<img alt="" src="https://assets.leetcode.com/uploads/2023/08/27/example2.png" style="width: 488px; height: 384px;" />
<pre>
<strong>Input:</strong> n = 6, edges = [[1,2],[1,3],[2,4],[3,5],[3,6]]
<strong>Output:</strong> 6
<strong>Explanation:</strong> 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.
</pre>

<p>&nbsp;</p>
<p><strong>Constraints:</strong></p>

<ul>
	<li><code>1 &lt;= n &lt;= 10<sup>5</sup></code></li>
	<li><code>edges.length == n - 1</code></li>
	<li><code>edges[i].length == 2</code></li>
	<li><code>1 &lt;= u<sub>i</sub>, v<sub>i</sub> &lt;= n</code></li>
	<li>The input is generated such that <code>edges</code> represent a valid tree.</li>
</ul>

# https://leetcode.com/problems/count-valid-paths-in-a-tree/solutions/4086958/dfs-dp-sieve-of-eratosthenes-beats-99-time-100-memory-with-clear-comments-in-code

def getAllPrimes():
    n = 10 ** 5
    prime = [True for _ in range(n + 1)]
    prime[0] = prime[1] = False
    x, y = 2, int(sqrt(n)) + 1
    while x < y:
        if prime[x]:
            for i in range(x * 2, n + 1, x):
                prime[i] = False
        x += 1
    return prime

class Solution:
    def countPaths(self, n: int, edges: List[List[int]]) -> int:
        primes = getAllPrimes()
        
        graph = [[] for _ in range(n + 1)]
        for src, dst in edges:
            graph[src].append(dst)
            graph[dst].append(src)
        
        res = 0
        dp = {}

        def dfs(root, par):
            nonlocal res
            allChain0s, allChain1s = [], []
            total0, total1 = 0, 0
            for child in graph[root]:
                if child == par: continue
                chain0, chain1 = dfs(child, root)
                total0 += chain0
                total1 += chain1
                allChain0s.append(chain0)
                allChain1s.append(chain1)
            if primes[root]:
                currChain0 = 0
                currChain1 = 1 + total0
                combTotal = 0
                for tmp0 in allChain0s:
                    combTotal += (tmp0 * (total0 - tmp0))
                res += total0 + combTotal // 2
            else:
                currChain0 = 1 + total0
                currChain1 = total1
                combTotal = 0
                for i in range(len(allChain0s)):
                    tmp0, tmp1 = allChain0s[i], allChain1s[i]
                    combTotal += (tmp0 * (total1 - tmp1)) + (tmp1 * (total0 - tmp0))
                res += total1 + combTotal // 2
            return currChain0, currChain1

        dfs(1, float("inf"))

        return res

    
90300