AtCoder が提供しているABC（AtCoder Beginner Contest）367 E問題をC++とPythonで解いてみました。ABC367は、2024年8月17日21:00に実施されました。
AtCoder の紹介はこちらに、プログラミングの方針はこちらに記事があります。
E問題 Permute K times（Difficulty : 1370）
問題の詳細は、リンク先をご覧ください。
ダブリングを適用して計算量を下げることができます。AtCoder Problems による Difficulty は 1370 でした。
解答案
繰り返す回数 $K$ の制約が $0 \leqq K \leqq 10^{18}$ と非常に大きく、定義通り計算していては間に合いません。
1回操作した後にもう1回操作をすると2回操作したことになります。2回操作した後にもう2回操作をすると4回操作したことになります。dp[i][j] を $2^i$ 回操作した後の要素 $j$ の場所とします。これは、以下の手順で求めることができます。
dp[i + 1][j] = dp[i][dp[i][j]]
十分な大きさの $i$ に対して、dp[i][j] を求めておくと、$K$ のセットされているビットについて演算します。例えば、11回後であれば、1回後→2回後→8回後のように場所を求めることができます。
C++ プログラム例（ABC367E）
上記の漸化式に従いdp[i][j] を計算します。dp[0][j] は、x[j] となります（21ー29行目）。
dp[i][j] を使って、k のビットに従い、行き先を求めます（33ー38行目）。最後に要素のインデックスから配列 a に変換して出力します。
以下が、C++プログラムです。
#include <bits/stdc++.h>
using namespace std;
typedef unsigned long long int ull;
int main()
{
int n;
ull k;
cin >> n >> k;
vector<ull> x(n);
for (int i = 0; i < n; ++i) {
cin >> x[i];
--x[i];
}
vector<ull> a(n);
for (int i = 0; i < n; ++i) {
cin >> a[i];
}
vector<vector<int>> dp(62, vector<int>(n, 0));
for (int i = 0; i < n; ++i) {
dp[0][i] = x[i];
}
for (int i = 0; i + 1 < 62; ++i) {
for (int j = 0; j < n; ++j) {
dp[i + 1][j] = dp[i][dp[i][j]];
}
}
vector<int> result(n);
for (int i = 0; i < n; ++i) {
result[i] = i;
for (int j = 0; j < 62; ++j) {
if ((k & (1ULL << j)) != 0) {
result[i] = dp[j][result[i]];
}
}
result[i] = a[result[i]];
}
for (int i = 0; i < n; ++i) {
cout << result[i] << " \n"[i == n - 1];
}
return 0;
}AC（Accepted＝正しいプログラム）と判定されました。
Python プログラム例（ABC367E）
Python版も基本的な考え方は同じです。
以下のプログラムは、「Python (CPython 3.11.4)」では、TLE判定となり、「Python (PyPy 3.10-v7.3.12)」で、AC判定となりました。
"""AtCoder Beginner Contest 367 E"""
n, k = map(int, input().split())
x = list(map(int, input().split()))
for i in range(n):
x[i] -= 1
a = list(map(int, input().split()))
dp = [[0 for _ in range(n)] for _ in range(62)]
for i in range(n):
dp[0][i] = x[i]
for i in range(61):
for j in range(n):
dp[i + 1][j] = dp[i][dp[i][j]]
result = [0] * n
for i in range(n):
result[i] = i
for j in range(62):
if (k & (1 << j)) != 0:
result[i] = dp[j][result[i]]
result[i] = a[result[i]]
print(*result)最後に
この問題は、ダブリングを使うことは分かりましたが、64ビットの整数を32ビットで演算していて、この間違いを時間内に気付くことができませんでした。残念です。
引き続き ABC の問題を紹介していきます。
