前提：記事中におけるグラフは、次の頂点 = f(現在の頂点) の形でグラフが定義される、いわゆる functional graph と呼ばれるものを想定しています。
大量のシミュレーションが必要で、愚直に実装すると TLE するけど、実はある部分はループ（閉路）を繰り返すだけで本当にシミュレーションが必要な部分は少ないという問題は時々出ます。
これを解くには、
- 閉路検出を実装する
- ダブリングを用いる
という方法がありますが、前者は手動で実装するには結構重いです。けんちょんさんの解説がわかりやすいのでソースコードをそのまま写させていただきますが、
#include <bits/stdc++.h>
using namespace std;
int main() {
int N;
long long K;
cin >> N >> K;
vector<int> A(N);
for (int i = 0; i < N; ++i) cin >> A[i], --A[i];
deque<int> a; // 繰り返しの部分
vector<bool> seen(N, false); // 一度見たかどうか
int cur = 0;
while (true) {
// 一度通った頂点を見つけたときの処理
if (seen[cur]) {
while (a[0] != cur) {
// 最初の余計な数手分を除去する
--K;
a.pop_front();
// 繰り返す前に K が限界になったらリターン
if (K == 0) {
cout << a[0]+1 << endl;
return 0;
}
}
break;
}
// 最初は愚直にシミュレーションしつつ、履歴をメモしていく
a.push_back(cur);
seen[cur] = true;
cur = A[cur];
}
cout << a[K % a.size()]+1 << endl;
}
メモ用の配列を別に用意したり、「最初の余計な数手分を除去する」「Kが限界になったらリターン」等の条件分岐が複雑です。本番でこういうのを一発で AC できる自身がない……。もっと逐次的に処理できないものか。
強連結成分
という訳で、強連結成分を検出するアルゴリズムである SCC を使ってこの問題を解いていきます。これを使うと、functional graph においては、ループ（閉路）を検出することができます。
なので、愚直なシミュレーションをほぼ流用する形で、
- 愚直なシミュレーションを続ける
- ループにたどり着いた段階で、残りのシミュレーション回数のうちループする分を一括処理する
- 愚直なシミュレーションを続ける
という逐次的な処理にできます。step2 が問題の本質ですね。
つまり、最終的には、その場所がループ上であるかを判別する（そうでない場合はループの長さを返す）ような配列 loop があるとして、
int pos = 0;
while (k--)
{
pos = a[pos];
if (loop[pos] > -1)
{
k %= loop[pos];
}
}
cout << pos + 1 << endl;
という形に持っていきたいです。こうすれば、kの場合分けなどを一切考えずに一括処理できます。
そのようなloopを用意するために、SCC を用いて以下のような下処理を行います。
ll n, k;
cin >> n >> k;
scc_graph graph(n); // SCCグラフの構造体
vector<int> a(n);
vector<ll> loop(n, -1); // ループの判別配列
rep(i, n)
{
cin >> a[i];
a[i]--;
graph.add_edge(i, a[i]);
if (i == a[i])
{
loop[i] = 1; // 自己ループの場合は個別処理
}
}
vector<vector<int>> scc = graph.scc();
for (auto v : scc)
{
if (v.size() == 1)
{
continue; // 強連結成分が 1 個の時は自己ループまたはループではない
}
for (int i : v)
{
loop[i] = v.size(); // 強連結成分が 2 個以上の時はループしている
}
}
ソースコード全体
#include <iostream>
using namespace std;
#define rep(i, n) for (int i = 0; i < n; i++)
#define rep1(i, n) for (int i = 1; i < n + 1; i++)
#define all(A) A.begin(), A.end()
#define itr(A, l, r) A.begin() + l, A.begin() + r
#define debug(var) cout << #var << " = " << var << endl;
typedef long long ll;
#include <atcoder/scc>
using namespace atcoder;
int main(void)
{
ll n, k;
cin >> n >> k;
scc_graph graph(n);
vector<int> a(n);
vector<ll> loop(n, -1);
rep(i, n)
{
cin >> a[i];
a[i]--;
graph.add_edge(i, a[i]);
if (i == a[i])
{
loop[i] = 1;
}
}
vector<vector<int>> scc = graph.scc();
for (auto v : scc)
{
if (v.size() == 1)
{
continue;
}
for (int i : v)
{
loop[i] = v.size();
}
}
int pos = 0;
while (k--)
{
pos = a[pos];
if (loop[pos] > -1)
{
k %= loop[pos];
}
}
cout << pos + 1 << endl;
}
注意点として、SCC は自己ループがあっても検出できないので、自己ループは個別に処理する必要があります。また、loop の初期値を 0 にすると、問題によっては 0 除算でエラーを起こす可能性があります（なので -1 にしています）。
もっと難しい問題
これはループを通る際に数値の加算があります。そのために配列を別に用意していますが、やることは基本的に同じです。
#include <iostream>
using namespace std;
#define rep(i, n) for (int i = 0; i < n; i++)
#define rep1(i, n) for (int i = 1; i < n + 1; i++)
#define all(A) A.begin(), A.end()
#define itr(A, l, r) A.begin() + l, A.begin() + r
#define debug(var) cout << #var << " = " << var << endl;
typedef long long ll;
#include <atcoder/scc>
using namespace atcoder;
int main(void)
{
int n;
ll k;
cin >> n >> k;
scc_graph graph(n);
vector<int> a(n);
vector<ll> loop(n, -1);
vector<ll> loop_plus(n);
rep(i, n)
{
int x;
cin >> x;
graph.add_edge(i, (i + x) % n);
a[i] = x;
if (i == (i + x) % n)
{
loop[i] = 1;
loop_plus[i] = x;
}
}
auto scc = graph.scc();
for (auto v : scc)
{
if (v.size() == 1)
{
continue;
}
ll plus = 0;
for (int i : v)
{
plus += a[i];
}
for (int i : v)
{
loop[i] = v.size();
loop_plus[i] = plus;
}
}
int pos = 0;
ll ans = 0;
while (k--)
{
ans += a[pos];
pos = ans % n;
if (loop[pos] > -1)
{
ans += loop_plus[pos] * (k / loop[pos]);
k %= loop[pos];
}
}
cout << ans << endl;
}