<p>There is a <strong>binary</strong> tree rooted at <code>0</code> consisting of <code>n</code> nodes. The nodes are labeled from <code>0</code> to <code>n - 1</code>. You are given a <strong>0-indexed</strong> integer array <code>parents</code> representing the tree, where <code>parents[i]</code> is the parent of node <code>i</code>. Since node <code>0</code> is the root, <code>parents[0] == -1</code>.</p>

<p>Each node has a <strong>score</strong>. To find the score of a node, consider if the node and the edges connected to it were <strong>removed</strong>. The tree would become one or more <strong>non-empty</strong> subtrees. The <strong>size</strong> of a subtree is the number of the nodes in it. The <strong>score</strong> of the node is the <strong>product of the sizes</strong> of all those subtrees.</p>

<p>Return <em>the <strong>number</strong> of nodes that have the <strong>highest score</strong></em>.</p>

<p>&nbsp;</p>
<p><strong class="example">Example 1:</strong></p>
<img alt="example-1" src="https://assets.leetcode.com/uploads/2021/10/03/example-1.png" style="width: 604px; height: 266px;" />
<pre>
<strong>Input:</strong> parents = [-1,2,0,2,0]
<strong>Output:</strong> 3
<strong>Explanation:</strong>
- The score of node 0 is: 3 * 1 = 3
- The score of node 1 is: 4 = 4
- The score of node 2 is: 1 * 1 * 2 = 2
- The score of node 3 is: 4 = 4
- The score of node 4 is: 4 = 4
The highest score is 4, and three nodes (node 1, node 3, and node 4) have the highest score.
</pre>

<p><strong class="example">Example 2:</strong></p>
<img alt="example-2" src="https://assets.leetcode.com/uploads/2021/10/03/example-2.png" style="width: 95px; height: 143px;" />
<pre>
<strong>Input:</strong> parents = [-1,2,0]
<strong>Output:</strong> 2
<strong>Explanation:</strong>
- The score of node 0 is: 2 = 2
- The score of node 1 is: 2 = 2
- The score of node 2 is: 1 * 1 = 1
The highest score is 2, and two nodes (node 0 and node 1) have the highest score.
</pre>

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

<ul>
	<li><code>n == parents.length</code></li>
	<li><code>2 &lt;= n &lt;= 10<sup>5</sup></code></li>
	<li><code>parents[0] == -1</code></li>
	<li><code>0 &lt;= parents[i] &lt;= n - 1</code> for <code>i != 0</code></li>
	<li><code>parents</code> represents a valid binary tree.</li>
</ul>

"""
思路： 用 post order DFS
因为 input 是 parent array，我们需要 convert 成 graph dict

1. 对于每个 node，我们都考虑 remove 该 node 后得到的 product 是多少
2. remove 掉 当前 node 后，整个 tree 会变成 3 parts
    a. left subtree: nodeCount is reported by left subtree
    b. right subtree: nodeCount is reported by right subtree
    c. parent tree: nodeCount = totalNodeCount - 1 - (left subtree node count + right subtree node count)

然后我们计算 remove 当前 node 后得到的 product

Time = O(V + E) due to build graph dictionary and DFS traversal
Space = O(V + E) due to graph dictionary
"""
from collections import defaultdict
class Solution:
    def countHighestScoreNodes(self, parents: List[int]) -> int:
        # edge case
        if not parents:
            return 0
        
        graph = self.__buildGraph(parents)
        self.productCount = defaultdict(int)
        totalNodeCount = len(parents)
        self.__dfs_traversal(0, graph, totalNodeCount)

        # find the max product
        maxProduct = max(self.productCount.keys())
        
        return self.productCount[maxProduct]
    
    def __dfs_traversal(self, curNode: int, graph: Dict[int, List[int]], totalNodeCount: int) -> None:
        totalSubtreeNodeCount = 0
        curProduct = 1
        for child in graph[curNode]:
            subtreeNodeCount = self.__dfs_traversal(child, graph, totalNodeCount)
            totalSubtreeNodeCount += subtreeNodeCount
            curProduct *= subtreeNodeCount
        
        """
        need to calculate the third subtree after removing current node
        minus 1 because we need to exlcude the current node as well

        Note: if currento node is root, totalNode - 1 - totalSubtreeNodeCount will be 0 
        so we need to use this max(1, ...) to ensure our product won't be zero
        """
        curProduct = curProduct * max(1, totalNodeCount - 1 - totalSubtreeNodeCount)
        self.productCount[curProduct] += 1

        # report the node count to parent
        return totalSubtreeNodeCount + 1

    def __buildGraph(self, parents: List[int]) -> Dict[int, List[int]]:
        graph = defaultdict(list)
        for index, parent in enumerate(parents):
            graph[parent].append(index)
        return graph

51212