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

class Solution:
    def countPaths(self, n: int, edges: List[List[int]]) -> int:
        is_prime = [True] * (n+1)
        is_prime[0] = False
        is_prime[1] = False
        for i in range(2, n+1):
            if is_prime[i]:
                for j in range(i+i, n+1, i):
                    is_prime[j] = False
        primes = set(i for i in range(n+1) if is_prime[i])

        # primes = [2]
        # for i in range(3, n+1, 2):
        #     isq = sqrt(i)
        #     is_prime = True
        #     for p in primes:
        #         if p > isq:
        #             break
        #         if i % p == 0:
        #             is_prime = False
        #             break
        #     if is_prime:
        #         primes.append(i)
        
        # primes = set(primes)

        g = defaultdict(list)
        for a, b in edges:
            g[a].append(b)
            g[b].append(a)

        mem = {}
        mem_fc = {}

        def num_zero_prime_paths(node, prev):
            if node in primes:
                return 0
            
            if (node, prev) in mem:
                return mem[(node, prev)]
            
            if node in mem_fc:
                # + 1 to account for the fact that this is not a prime
                return mem_fc[node] - num_zero_prime_paths(prev, node) + 1
                        
            result = 1

            for c in g[node]:
                if c == prev:
                    continue
                
                result += num_zero_prime_paths(c, node)

            mem[(node, prev)] = result
            return result

        nc = defaultdict(int)
        for a, b in edges:
            nc[a] += 1
            nc[b] += 1
        
        for _, node in sorted(((c, node) for node, c in nc.items()), reverse=True):
            fc = 0
            for c in g[node]:
                fc += num_zero_prime_paths(c, node)
            mem_fc[node] = fc
        
        
        
        res = 0

        for p in primes:
            zero_prime_branch_counts = [
                num_zero_prime_paths(c, p)
                for c in g[p]
            ]
            # print(f"p {p}, our zero prime branches are {zero_prime_branch_counts}")
            zero_prime_branch_counts = [e for e in zero_prime_branch_counts if e > 0]
            if len(zero_prime_branch_counts) == 0:
                continue
            ps = [0]
            for cnt in zero_prime_branch_counts:
                ps.append(ps[-1] + cnt)
            
            for i in range(len(zero_prime_branch_counts)-1, -1, -1):
                # +1 to account for paths starting at the prime
                res += zero_prime_branch_counts[i] * (ps[i] + 1)
            
        return res
1831