📑 Chain Reactions

Codejam Qualification Round 2022. Task 4

Python Only

class DGraph:
    def __init__(self,n,f,p):
        self.N=n
        self.P=p
        self.F=f
        self.G={k:0 for k in range(n)}
        self.D={k:0 for k in range(n)}
    def add_edges(self):
        for i,el in enumerate(self.P):
            self.G[i]=el-1
            if el!=0:
                self.D[el-1]+=1 
    def max_fun(self):
        indegree=self.D
        inits=[]
        for i in range(self.N):
            if indegree[i]==0:
                inits.append(i)
        fmin=[10**12]*self.N
        fsum=sum([self.F[i] for i in inits])
        while inits:
            i=inits.pop(0)
            node=self.G[i]
            if node!=-1:
                indegree[node]-=1
                fmin[node]=min([fmin[node],self.F[i]])
                if indegree[node]==0:
                    inits.append(node)
                    self.F[node]=max([self.F[node],fmin[node]])
                    fsum+=self.F[node]-fmin[node]
        return fsum
T=int(input())
for t in range(T):
    N=int(input())
    F=list(map(int,input().split()))
    P=list(map(int,input().split()))
    dgraph=DGraph(N,F,P)
    dgraph.add_edges()
    FSUM=dgraph.max_fun()
    print('Case #{}: {}'.format(t+1,FSUM))