ABC337-C 解説

C - Lining Up 2

解法1

人 $i$ に対してその後ろにいる人 $B_i$ を求める方法です。$B$ を求める時に $A_i=-1$ となる $i$ に対しては列の先頭なのでそれを持つ必要もあります。

N = int(input())
A = list(map(int, input().split()))
B = [-1] * N
now = 0
for i in range(N):
    if A[i] == -1:
        now = i
    else:
        B[A[i] - 1] = i
ans = [now + 1]
for i in range(N - 1):
    now = B[now]
    ans.append(now + 1)
print(*ans)

#include <bits/stdc++.h>
using namespace std;

int main() {
    int N;
    cin >> N;
    vector<int> A(N);
    for (int i = 0; i < N; i++) {
        cin >> A[i];
    }
    vector<int> B(N, -1);
    int now = 0;
    for (int i = 0; i < N; i++) {
        if (A[i] == -1) {
            now = i;
        } else {
            B[A[i] - 1] = i;
        }
    }
    vector<int> ans = {now + 1};
    for (int i = 0; i < N - 1; i++) {
        now = B[now];
        ans.push_back(now + 1);
    }
    for (int i = 0; i < N; i++) {
        cout << ans[i] << ' ';
    }
}

解法2

列を後ろから決定していく方法です。最初に $A$ に現れない値を求めてさかのぼっていきます。最後に逆順に出力することに注意してください。

N = int(input())
A = list(map(int, input().split()))
cnt = [0] * N
for i in A:
    if i != -1:
        cnt[i - 1] = 1
now = -1
for i in range(N):
    if cnt[i] == 0:
        now = i
ans = [now + 1]
for _ in range(N - 1):
    now = A[now] - 1
    ans.append(now + 1)
print(*ans[::-1])

#include <bits/stdc++.h>
using namespace std;

int main() {
    int N;
    cin >> N;
    vector<int> A(N);
    for (int i = 0; i < N; i++) {
        cin >> A[i];
    }
    vector<int> cnt(N);
    for (int i: A) {
        if (i != -1) {
            cnt[i - 1] = 1;
        }
    }
    int now = -1;
    for (int i = 0; i < N; i++) {
        if (cnt[i] == 0) {
            now = i;
        }
    }
    vector<int> ans = {now + 1};
    for (int i = 0; i < N - 1; i++) {
        now = A[now] - 1;
        ans.push_back(now + 1);
    }
    for (int i = N - 1; i >= 0; i--) {
        cout << ans[i] << ' ';
    }
}