Tree game

import heapq
import sys
input = sys.stdin.readline
sys.setrecursionlimit(100002)

def topo_sort():
    order = []
    visited = [0]*(n+1)
    def dfs(u):
        if u==1:
            visited[1]=True
            return True
        visited[u]=-1
        for v in down_graph[u]:
            # v = v[0]
            if visited[v] ==-1:
                return False 
            if visited[v]==0 and not dfs(v):
                return False
        visited[u] = True
        order.append(u)
        return True

    
    for node in range(1,n+1):
        if visited[node] ==0 and not dfs(node):
            print("IMPOSSIBLE")
            return []
    # return order[::-1]
    order.append(1)
    return order




n = int(input())
ai  = list(map(int, input().split()))
pi  = list(map(int, input().split()))
if n ==1:
    print(ai[0])
    exit()
nodes_sorted_by_down_neighbours = []
down_graph = {j:[] for j in range(1,n+1)}
dp = []
for j in range(n):
    heapq.heappush(down_graph[pi[j]], j+1) 

# print(up_graph, down_graph)

scores = [0]*(n+1)
order = topo_sort()
for node in order:
    scores[node] = ai[node-1]
    k = 3 if node == 1 else 2
    next_scores = [scores[i] for i in down_graph[node] if i!=1]
    if len(next_scores)>= k :
        next_scores.sort()
        scores[node] += next_scores[-k]
        # print(next_scores)
# print(scores)
print(scores[1])

Test 1

Verdict: ACCEPTED

Test 2

Verdict: ACCEPTED

Test 3

Verdict: ACCEPTED

Test 4

Verdict: ACCEPTED

Test 5

Verdict: WRONG ANSWER

Test 6

Verdict: ACCEPTED

Test 7

Verdict: WRONG ANSWER

Test 8

Verdict: WRONG ANSWER

Test 9

Verdict: WRONG ANSWER

Test 10

Verdict: ACCEPTED

Test 11

Verdict: WRONG ANSWER

Test 12

Verdict: ACCEPTED

Test 13

Verdict: WRONG ANSWER

Test 14

Verdict: ACCEPTED

Test 15

Verdict: WRONG ANSWER

Test 16

Verdict: WRONG ANSWER

Test 17

Verdict: WRONG ANSWER

Test 18

Verdict: ACCEPTED

Test 19

Verdict: WRONG ANSWER

Test 20

Verdict: ACCEPTED

Test 21

Verdict: WRONG ANSWER

Test 22

Verdict: WRONG ANSWER

Test 23

Verdict: ACCEPTED

Test 24

Verdict: WRONG ANSWER

Test 25

Verdict: WRONG ANSWER

Test 26

Verdict: WRONG ANSWER

Test 27

Verdict: ACCEPTED

Test 28

Verdict: WRONG ANSWER

Test 29

Verdict: WRONG ANSWER

Test 30

Verdict: WRONG ANSWER

Test 31

Verdict: WRONG ANSWER

Test 32

Verdict: WRONG ANSWER

Test 33

Verdict: WRONG ANSWER

Test 34

Verdict: WRONG ANSWER

Test 35

Verdict: WRONG ANSWER

Test 36

Verdict: WRONG ANSWER

Test 37

Verdict: WRONG ANSWER

Test 38

Verdict: ACCEPTED

Test 39

Verdict: WRONG ANSWER

Test 40

Verdict: WRONG ANSWER

Test 41

Verdict: WRONG ANSWER

Test 42

Verdict: WRONG ANSWER

Test 43

Verdict: WRONG ANSWER

Test 44

Verdict: WRONG ANSWER

Test 45

Verdict: WRONG ANSWER

Test 46

Verdict: WRONG ANSWER

Test 47

Verdict: WRONG ANSWER

Test 48

Verdict: WRONG ANSWER

Test 49

Verdict: WRONG ANSWER

Test 50

Verdict: WRONG ANSWER

Test 51

Verdict: ACCEPTED

Test 52

Verdict: WRONG ANSWER

Test 53

Verdict: WRONG ANSWER

Test 54

Verdict: WRONG ANSWER

Test 55

Verdict: WRONG ANSWER

Test 56

Verdict: WRONG ANSWER

Test 57

Verdict: WRONG ANSWER

Test 58

Verdict: ACCEPTED

Test 59

Verdict: WRONG ANSWER

Test 60

Verdict: WRONG ANSWER

Test 61

Verdict: WRONG ANSWER

Test 62

Verdict: WRONG ANSWER

Test 63

Verdict: WRONG ANSWER

Test 64

Verdict: ACCEPTED

Test 65

Verdict: WRONG ANSWER

Test 66

Verdict: WRONG ANSWER

Test 67

Verdict: WRONG ANSWER

Test 68

Verdict: WRONG ANSWER

Test 69

Verdict: WRONG ANSWER