Problem p02809 (python)

Fixed submission:
import sys
from collections import defaultdict
from heapq import heappop, heappush
from itertools import permutations
from operator import itemgetter


# ある条件に当てはまらない限り、i番目に置く数字 xi は、
# それまで使ってない中で最も小さい数字か、
# またはその次に小さい数字（x[i-1]の右に最も小さい数字を置けない場合）
#
# ある条件: 以下の条件を満たす、未使用の数 k がある
# 残っているk以外の全ての数字が、kを共通して右側に置けない数として指定している
#   ＝kを先頭に持ってこない限り、kを置ける機会が無い
#
# ただし残りが少なく（3以下）なってくると例外的なものが出てくるため、それ以降は全探索

def fill_remainings(ans, aaa, x, remainings):
    """
    xを先頭にして残りを昇順に追加
    ただしxの次の要素のみ、aaa[x]で禁止されていた場合はその次と入れ替える
    remainingsにはxを含め3要素以上残っていることが前提
    """
    ans.append(x)
    i = len(ans)
    while remainings:
        k = heappop(remainings)
        if k != x:
            ans.append(k)
    if aaa[x] == ans[i]:
        ans[i], ans[i + 1] = ans[i + 1], ans[i]


def solve(n, aaa):
    if n == 2:
        return [-1]

    in_degrees = defaultdict(lambda: 0)
    for i, a in enumerate(aaa, start=1):
        in_degrees[a] += 1
    in_degrees = dict(in_degrees)

    # 少なくとも残り個数がこれ+1になるまでは「ある条件」には当てはまらない
    # ただし減少することはあるため、直前に再チェック必要
    curr_max = max(in_degrees.values())

    remainings = list(range(1, n + 1))

    aaa.insert(0, 0)

    ans = []
    banned = -1
    for i in range(n - 3):
        if curr_max == n - i - 1:
            curr_x, curr_max = max(in_degrees.items(), key=itemgetter(1))
            if curr_max == n - i - 1:
                fill_remainings(ans, aaa, curr_x, remainings)
                return ans
        top = heappop(remainings)
        if top == banned:
            ans.append(heappop(remainings))
            heappush(remainings, top)
        else:
            ans.append(top)
        banned = aaa[ans[-1]]
        # 確定した数字の入り次数を削減
        if banned in in_degrees:
            if in_degrees[banned] == 1:
                del in_degrees[banned]
            else:
                in_degrees[banned] -= 1
        in_degrees.pop(ans[-1], 0)

    remainings.sort()
    for i, j, k in permutations(remainings):
        if i != banned and j != aaa[i] and k != aaa[j]:
            ans += [i, j, k]
            break

    return ans


n, *aaa = map(int, sys.stdin.buffer.read().split())
print(*solve(n, aaa))

Buggy submission (verdict: Runtime Error):
import sys
from collections import defaultdict
from heapq import heappop, heappush
from itertools import permutations
from operator import itemgetter


# ある条件に当てはまらない限り、i番目に置く数字 xi は、
# それまで使ってない中で最も小さい数字か、
# またはその次に小さい数字（x[i-1]の右に最も小さい数字を置けない場合）
#
# ある条件: 以下の条件を満たす、未使用の数 k がある
# 残っているk以外の全ての数字が、kを共通して右側に置けない数として指定している
#   ＝kを先頭に持ってこない限り、kを置ける機会が無い
#
# ただし残りが少なく（3以下）なってくると例外的なものが出てくるため、それ以降は全探索

def fill_remainings(ans, aaa, x, remainings):
    """
    xを先頭にして残りを昇順に追加
    ただしxの次の要素のみ、aaa[x]で禁止されていた場合はその次と入れ替える
    remainingsにはxを含め3要素以上残っていることが前提
    """
    ans.append(x)
    i = len(ans)
    while remainings:
        k = heappop(remainings)
        if k != x:
            ans.append(k)
    if aaa[x] == ans[i]:
        ans[i], ans[i + 1] = ans[i + 1], ans[i]


def solve(n, aaa):
    if n == 2:
        return [-1]

    in_degrees = defaultdict(lambda: 0)
    for i, a in enumerate(aaa, start=1):
        in_degrees[a] += 1
    in_degrees = dict(in_degrees)

    # 少なくとも残り個数がこれ+1になるまでは「ある条件」には当てはまらない
    # ただし減少することはあるため、直前に再チェック必要
    curr_max = max(in_degrees.values())

    remainings = list(range(1, n + 1))

    aaa.insert(0, 0)

    ans = []
    banned = -1
    for i in range(n - 3):
        if curr_max == n - i - 1:
            curr_x, curr_max = max(in_degrees.items(), key=itemgetter(1))
            if curr_max == n - i - 1:
                fill_remainings(ans, aaa, curr_x, remainings)
                return ans
        top = heappop(remainings)
        if top == banned:
            ans.append(heappop(remainings))
            heappush(remainings, top)
        else:
            ans.append(top)
        banned = aaa[ans[-1]]
        # 確定した数字の入り次数を削減
        if in_degrees[banned] == 1:
            del in_degrees[banned]
        else:
            in_degrees[banned] -= 1
        in_degrees.pop(ans[-1], 0)

    remainings.sort()
    for i, j, k in permutations(remainings):
        if i != banned and j != aaa[i] and k != aaa[j]:
            ans += [i, j, k]
            break

    return ans


n, *aaa = map(int, sys.stdin.buffer.read().split())
print(*solve(n, aaa))