競技プログラミング初心者のためのatCoder参加記録（Beginner Contest 337）

問題文

2024年1月20日（土）21時に行われた ABC337 の C問題について解説します。

問題解説

問題文の入力例１を例にとって確認します。

入力値Nが６なので、１列に６人並んでいます。先頭から人1, 人2, 人3, 人4, 人5, 人6と呼びます。

図1 人1～6

次の入力値A1～A6（4 1 -1 5 3 2）は、人nがどの人の後ろに並んでいるのかを示しています。-1 は先頭です。図で表すと以下のようになります。この値が後ろの人が誰なのかを示すポインターとして使用できます。

図2 A1～6

そのポインターをキーとして順番に並び替えた図が以下になります。

図3 並び替え

そして、先頭から順にその人の番号を出力すればOKですね。

解答例

図２を参考に入力値は、以下のような体系で取り込みます。

入力値をどう扱うかを考えます。制約条件で人の数は最大 3x10⁵ なので、ポインターをキーとしてForループで後ろの人を探していたら TLEに引っかかりそうです。ですので、ポインターをキー値としてもつ辞書型を使用してこの問題をクリアします。

図４　入力値を辞書型で取り込む

これをコーディングしたものが以下です。enumerate関数を使用して取得したインデックスは、後ろの人が誰なのかを示すvalue値として使用します。

N=int(input())
#> N：6

#入力値Aを取り込みます。
A=list(map(int, input().split()))
#> A：[4, 1, -1, 5, 3, 2]

#入力値を辞書型で扱えるようにする。
dicA={}
for i, a in enumerate(A):
    dicA[a]=i+1
#> dicA：{4: 1, 1: 2, -1: 3, 5: 4, 3: 5, 2: 6}

「#入力値を辞書型で扱えるようにする。」部分は以下のように１行にすることが可能です。

dicA={a:i+1 for i, a in enumerate(A)}

そして、次に先頭から順に後ろの人を探していく部分をコーディングします。処理を図で示すと以下のようになります。

図５　後ろの人を探していく処理

コーディングすると以下のようになります。

ans=[]  #答えを格納するリストを初期化
key=-1  #キー値を格納する変数を-1：先頭にしておく。
#人数分ループ
for _ in range(N):
    #キー値を元にうしろの人のキー値（value）を取得する。
    key=dicA[key]
    #答えリストにうしろの人のキー値を追加
    ans.append(key)
#> ans：[3, 5, 4, 1, 2, 6]

#答えリストを出力
print(*ans)
#> 3 5 4 1 2 6

最後に

辞書型を使用することで、見た目もスッキリして TLEもクリアできるプログラムが作成できました。