Google Code Jam 2022 -- Qualification Round

chain_reactions.py

import sys
from operator import iconcat
from functools import reduce

sys.setrecursionlimit(100010)


class Node(object):
    def __init__(self):
        self.fun = 0
        self.children = []

    def __repr__(self):
        if not self.children:
            return f'Leaf({self.fun})'
        return f'Node({self.fun}, {self.children})'

    def reduce_tree(self):
        while len(self.children) == 1:
            child = self.children[0]
            self.fun = max(self.fun, child.fun)
            self.children = child.children
        for child in self.children:
            child.reduce_tree()

    def aux(self):
        if not self.children:
            return (self.fun, 0)
        candidates = [child.aux() for child in self.children]
        most_serious = min(candidates)[0]
        accumulate = sum(map(sum, candidates)) - most_serious
        representator = max(most_serious, self.fun)
        return (representator, accumulate)

    def maximum_fun(self):
        self.reduce_tree()
        return sum(self.aux())


def main():
    for t in range(int(input())):
        n = int(input())
        nodes = [Node() for _ in range(n+1)]
        for i, fun in enumerate(map(int, input().split()), 1):
            nodes[i].fun = fun
        for child, parent in enumerate(map(int, input().split()), 1):
            nodes[parent].children += [nodes[child]]
        answer = nodes[0].maximum_fun()
        print(f'Case #{t+1}: {answer}')


if __name__ == '__main__':
    main()