Captain Flint and Treasure (Codeforces 1388D)
Tags: Greedy, trees

Captain Fint is involved in another treasure hunt, but have found only one strange problem. The problem may be connected to the treasure's location or may not. That's why captain Flint decided to leave the solving the problem to his crew and offered an absurdly high reward: one day off. The problem itself sounds like this...There are two arrays $$$a$$$ and $$$b$$$ of length $$$n$$$. Initially, an $$$ans$$$ is equal to $$$0$$$ and the following operation is defined:   Choose position $$$i$$$ ($$$1 \le i \le n$$$);  Add $$$a_i$$$ to $$$ans$$$;  If $$$b_i \neq -1$$$ then add $$$a_i$$$ to $$$a_{b_i}$$$. What is the maximum $$$ans$$$ you can get by performing the operation on each $$$i$$$ ($$$1 \le i \le n$$$) exactly once?Uncle Bogdan is eager to get the reward, so he is asking your help to find the optimal order of positions to perform the operation on them.

Input: The first line contains the integer $$$n$$$ ($$$1 \le n \le 2 \cdot 10^5$$$) — the length of arrays $$$a$$$ and $$$b$$$.The second line contains $$$n$$$ integers $$$a_1, a_2, \ldots, a_n$$$ ($$$−10^6 \le a_i \le 10^6$$$).The third line contains $$$n$$$ integers $$$b_1, b_2, \ldots, b_n$$$ ($$$1 \le b_i \le n$$$ or $$$b_i = -1$$$).Additional constraint: it's guaranteed that for any $$$i$$$ ($$$1 \le i \le n$$$) the sequence $$$b_i, b_{b_i}, b_{b_{b_i}}, \ldots$$$ is not cyclic, in other words it will always end with $$$-1$$$.

Output: In the first line, print the maximum $$$ans$$$ you can get.In the second line, print the order of operations: $$$n$$$ different integers $$$p_1, p_2, \ldots, p_n$$$ ($$$1 \le p_i \le n$$$). The $$$p_i$$$ is the position which should be chosen at the $$$i$$$-th step. If there are multiple orders, print any of them.


Sample Inputs:
3
1 2 3
2 3 -1

2
-1 100
2 -1

10
-10 -1 2 2 5 -2 -3 -4 2 -6
-1 -1 2 2 -1 5 5 7 7 9

Sample Outputs:
10
1 2 3

99
2 1

-9
3 5 6 1 9 4 10 7 8 2

Solution (Python):
import sys
input = sys.stdin.readline

def readList():
    return list(map(int, input().split()))

def readInt():
    return int(input())

def readInts():
    return map(int, input().split())

def readStr():
    return input().strip()

def solve():
    (n, A, B) = (readInt(), readList(), readList())
    graph = [[] for _ in range(n)]
    par = [-1] * n
    for i in range(n):
        if B[i] != -1:
            graph[B[i] - 1].append(i)
            par[i] = B[i] - 1
    (now, after) = ([], [])
    for i in range(n):
        if B[i] == -1:
            st = [(i, 0)]
            while st:
                (node, exp) = st.pop()
                if exp:
                    if node == i:
                        now.append(node + 1)
                        continue
                    if A[node] >= 0:
                        A[par[node]] += A[node]
                        now.append(node + 1)
                    else:
                        after.append(node + 1)
                else:
                    st.append((node, 1))
                    for nei in graph[node]:
                        st.append((nei, 0))
    print(sum(A))
    print(*now + after[::-1])
    return
solve()

Solution (Python):
from sys import stdin, stdout

def captain_flint_and_treasure(n, a_a, b_a):
    ind = [0 for i in range(n + 1)]
    r_s = set((i + 1 for i in range(n)))
    for i in range(n):
        if b_a[i] != -1:
            ind[b_a[i]] += 1
            r_s.discard(b_a[i])
    queue = list(r_s)
    ans1 = 0
    ans2 = []
    ans3 = []
    while queue:
        cnt = len(queue)
        for i in range(cnt):
            j = queue.pop()
            if b_a[j - 1] != -1:
                if a_a[j - 1] > 0:
                    a_a[b_a[j - 1] - 1] += a_a[j - 1]
                ind[b_a[j - 1]] -= 1
                if ind[b_a[j - 1]] == 0:
                    queue.append(b_a[j - 1])
            ans1 += a_a[j - 1]
            if a_a[j - 1] > 0:
                ans2.append(j)
            else:
                ans3.append(j)
    ans3.reverse()
    return [ans1, ans2 + ans3]
n = int(stdin.readline())
a_a = list(map(int, stdin.readline().split()))
b_a = list(map(int, stdin.readline().split()))
ans = captain_flint_and_treasure(n, a_a, b_a)
stdout.write(str(ans[0]) + '\n')
stdout.write(' '.join(map(str, ans[1])) + '\n')

Solution (Python):
from collections import defaultdict
import sys
sys.setrecursionlimit(2 * 10 ** 5)
n = int(input())
a = [0] + list(map(int, input().split()))
b = list(map(int, input().split()))
f = defaultdict(list)
num_list = [0] * (n + 1)
for (index, nxt) in enumerate(b):
    if nxt != -1:
        f[index + 1].append(nxt)
        num_list[nxt] += 1
q = [i for (i, v) in enumerate(num_list) if i > 0 and v == 0]
res = 0
res_list = []
reverse_nege_list = []
while q:
    p = []
    for i in q:
        if a[i] > 0:
            res_list.append(i)
        else:
            reverse_nege_list.append(i)
        res += a[i]
        for nxt in f[i]:
            num_list[nxt] -= 1
            if num_list[nxt] == 0:
                p.append(nxt)
            if a[i] > 0:
                a[nxt] += a[i]
    q = p
print(res)
ans_list = res_list + reverse_nege_list[::-1]
print(' '.join(map(str, ans_list)))