Count Nodes With the Highest Score LeetCode solution in Python

1
class Solution:
2
    def countHighestScoreNodes(self, parents: List[int]) -> int:
3
        graph = collections.defaultdict(list)
4
        for node, parent in enumerate(parents):  # build graph
5
            graph[parent].append(node)
6
        n = len(parents)  # total number of nodes
7
        d = collections.Counter()
8

9
        def count_nodes(node):  # number of children node + self
10
            p, s = 1, 0  # p: product, s: sum
11
            for child in graph[node]:  # for each child (only 2 at maximum)
12
                res = count_nodes(child)  # get its nodes count
13
                p *= res  # take the product
14
                s += res  # take the sum
15
            p *= max(
16
                1, n - 1 - s
17
            )  # times up-branch (number of nodes other than left, right children ans itself)
18
            d[p] += 1  # count the product
19
            return s + 1  # return number of children node + 1 (self)
20

21
        count_nodes(0)  # starting from root (0)
22
        return d[max(d.keys())]  # return max count