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

<p>A node is <strong>good</strong> if all the <span data-keyword="subtree">subtrees</span> rooted at its children have the same size.</p>

<p>Return the number of <strong>good</strong> nodes in the given tree.</p>

<p>A <strong>subtree</strong> of <code>treeName</code> is a tree consisting of a node in <code>treeName</code> and all of its descendants.</p>

<p>&nbsp;</p>
<p><strong class="example">Example 1:</strong></p>

<div class="example-block">
<p><strong>Input:</strong> <span class="example-io">edges = [[0,1],[0,2],[1,3],[1,4],[2,5],[2,6]]</span></p>

<p><strong>Output:</strong> <span class="example-io">7</span></p>

<p><strong>Explanation:</strong></p>
<img alt="" src="https://assets.leetcode.com/uploads/2024/05/26/tree1.png" style="width: 360px; height: 158px;" />
<p>All of the nodes of the given tree are good.</p>
</div>

<p><strong class="example">Example 2:</strong></p>

<div class="example-block">
<p><strong>Input:</strong> <span class="example-io">edges = [[0,1],[1,2],[2,3],[3,4],[0,5],[1,6],[2,7],[3,8]]</span></p>

<p><strong>Output:</strong> <span class="example-io">6</span></p>

<p><strong>Explanation:</strong></p>
<img alt="" src="https://assets.leetcode.com/uploads/2024/06/03/screenshot-2024-06-03-193552.png" style="width: 360px; height: 303px;" />
<p>There are 6 good nodes in the given tree. They are colored in the image above.</p>

<p><strong class="example">Example 3:</strong></p>

<div class="example-block">
<p><strong>Input:</strong> <span class="example-io">edges = [[0,1],[1,2],[1,3],[1,4],[0,5],[5,6],[6,7],[7,8],[0,9],[9,10],[9,12],[10,11]]</span></p>

<p><strong>Output:</strong> <span class="example-io">12</span></p>

<p><strong>Explanation:</strong></p>
<img alt="" src="https://assets.leetcode.com/uploads/2024/08/08/rob.jpg" style="width: 450px; height: 277px;" />
<p>All nodes except node 9 are good.</p>
</div>
</div>

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

<ul>
	<li><code>2 &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>0 &lt;= a<sub>i</sub>, b<sub>i</sub> &lt; n</code></li>
	<li>The input is generated such that <code>edges</code> represents a valid tree.</li>
</ul>

class Solution:
    def countGoodNodes(self, edges: List[List[int]]) -> int:
        mapping = {}
        n = 0
        for k, v in edges:
            mapping.setdefault(k, set()).add(v)
            mapping.setdefault(v, set()).add(k)
            n = max(n, k, v)
        
        seen = set()
        stack = [0]
        # while stack:
        #     new_stack = []
        #     for node in stack:
        #         mapping[node].difference_update(seen)
        #         new_stack.extend(mapping[node])
        #         seen.add(node)
        for node in stack:
            mapping[node].difference_update(seen)
            stack.extend(mapping[node])
            seen.add(node)

            
            # stack = new_stack
        
        # return mapping
        

        # print(f"{mapping = }")      
        # print(f"{linkages = }")

        # return mapping
        
        self.mapping = mapping
        self.memo = {}

        # get lengths of each node
        self.recur(0)
        # print(f"{self.memo = }")
        
        # iterate thru memo
        self.count_memo = {}
        total = self.good_nodes(0)
    
        return total
    
    
    def recur(self, curr):
        # base case
        # if curr not in self.mapping:
        #     self.memo[curr] = 0
        #     return self.memo[curr]
        if not self.mapping.get(curr):
            self.memo[curr] = 0
            return self.memo[curr]
        if curr in self.memo:
            return self.memo[curr]
        
        # main logic
        lengths = []
        count = 0
        for child in self.mapping[curr]:
            lengths.append(self.recur(child))
            count += 1

        self.memo[curr] = sum(lengths) + count
        # print(f"{curr, self.memo[curr] = }")
        return self.memo[curr]
    
    def good_nodes(self, curr):
        # base case
        if curr in self.count_memo:
            return self.count_memo[curr]
        if self.memo[curr] == 0:
            return 1
        
        # main logic
        lengths = []
        for child in self.mapping[curr]:
            lengths.append(self.memo[child])
        
        equal_length = lengths[0]        
        is_good_node = sum([1 for length in lengths if length == equal_length]) == len(lengths)
        
        ans = 0
        for child in self.mapping[curr]:
            ans += self.good_nodes(child)
        
        self.count_memo[curr] = ans + (1 if is_good_node else 0)
        return self.count_memo[curr]
120481