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

'''
    1
   2 3
  4 5  
      
      
What I get from the other side matters too -> in trees always consider this
Hard part of this problem is understanding that you can go from any point in one tree to another part of the subtree -> so like considering all pairs of paths

Seems like I need to find the "root" because if you think about it -> a V shaped, my algo won't work so choosing the right starting point is important, you basically can't have a node in the same "level" otherwise it won't work    
'''

class Solution:
    def countPaths(self, n: int, edges: List[List[int]]) -> int:
        neighs = defaultdict(list)
        visited = set()
        res = [0]
        
        p = 2
        primes = [True] * (n + 1)
        primes[1] = False
        
        while p * p <= n:
            if primes[p]:
                for i in range(p * p, n + 1, p):
                    primes[i] = False
            p += 1
        
        for edge in edges:
            neighs[edge[0]].append(edge[1])
            neighs[edge[1]].append(edge[0])
            
        self.countPathsHelper(1, primes, neighs, res, visited, 0, 0)
        return res[0]
            
    def countPathsHelper(self, n: int, primes, neighs, res, visited, no_prime, one_prime):
        visited.add(n)
        tot_non_prime_paths = 0
        tot_one_prime_paths = 0
        is_prime = primes[n]
  
        for neigh in neighs[n]:
            if neigh not in visited:
                visited.add(neigh)
                # Need to comment below -> a bit tricky
                # Need to pass in the other side of the tree's results
                # The missing bit here is that you need to pass what you know above to others...
                # Meaning the no_prime + tot_non_prime_paths (key point here)
                # Kind of tricky with the other_prime because it's tot_non_prime_paths + no_prime which is no symmetrical
                # Cool tree problemw here you need to do everything...pass from above, what do I return to my parent, pass from the side
                other_non_prime = 0 if is_prime else no_prime + tot_non_prime_paths
                other_prime = tot_non_prime_paths + no_prime if is_prime else one_prime + tot_one_prime_paths
                a, b = self.countPathsHelper(neigh, primes, neighs, res, visited, other_non_prime, other_prime)
                tot_non_prime_paths += a
                tot_one_prime_paths += b
            
        if is_prime:
            res[0] += tot_non_prime_paths + no_prime          
            return 0, tot_non_prime_paths + 1
        else:
            res[0] += tot_one_prime_paths + one_prime
            return tot_non_prime_paths + 1, tot_one_prime_paths
  
            
        
        
        
1436