集合 \(S=\{ 0,1,\ldots, N-1 \}\) に対して, 関数 \(f:S\to S\) を \(f(i)=(i+A_i)\bmod N\) で定めます。ここで \(\mod N\) は問題文中同様、 \(N\) で割った余りをさします。任意の非負整数 \(k\) について、\(k\) 回操作した後のアメの個数を \(N\) で割った余りは \(f^k(0)\) です。ここで \(f^k\) は \(f\) を \(k\) 回合成したものをさします。ただし、\(f^0(x)=x\) です。求めたいものは \(A_0+A_{f(0)}+\cdots+A_{f^{K-1}(0)}\) です。また、\(f^k(0)=\bigl( A_0+A_{f(0)}+\cdots +A_{f^{k-1}(0)}\bigr)\% N\) とも書けることに注意します。

解法1 : 周期を使う方法

\(f^0(0),f^1(0),\ldots,f^N(0)\) を考えるとこれらはすべて \(0\) 以上 \(n\) 未満の整数であるので、ある異なる \(0\leq k<k' \leq N\)が存在して \(f^k(0)=f^{k'}(0)\) となります。ここで、ある \(0\leq s<t\) が存在して、\(f^s(0)=f^t(0)\) となる最小の \(t\) を取り、\(p=t-s\)とします。このとき、任意の非負整数 \(i\)に対して, \(f^{t+i}(0)=f^i(f^t(0))=f^i(f^s(0))=f^{s+i}(0)\) となります。つまり、\(s\leq k \)なる任意の\(k\) について \(f^{k+p}(0)=f^k(0)\)が成り立ち、数列 \(\{ f^i(0) \}\) は\(s\leq i\) の範囲で周期 \(p\) を持つ事が分かります。これを用いて答えを求めます。

\(K\leq s\) のときは、\(K\leq s<t\leq N\) であるので、愚直にアメの個数の遷移をシミュレーションすることができます。そうでないとき、非負整数 \(a\) と \(p\) 未満の非負整数 \(b\) を用いて \((K-1-s)=ap+b\) と書けます。 ここで、非負整数 \(i\) に対して、

\[ A_{f^{s+ip+b+1}(0)}+A_{f^{s+ip+b+2}(0)}+\cdots+A_{f^{s+(i+1)p+b}(0)}= A_{f^{s+b+1}(0)}+A_{f^{s+b+2}(0)}+\cdots+A_{f^{t-1}(0)}+A_{f^s(0)}+A_{ f^{s+1}(0)}+\cdots+A_{f^{s+b}(0)} \} \]

が成り立ち, \(X=A_{f^s(0)}+A_{ f^{s+1}(0)}+\cdots+A_{f^{t-1}(0)}\) とすると, \(A_0+A_{f(0)}+\cdots+A_{f^{K-1}(0)}= (A_0+A_{f(0)}+\cdots+A_{f^{s+b}(0)})+(A_{f^{s+b+1}(0)}+A_{f^{s+b+2}(0)}+\cdots+A_{f^{s+p+b}(0)})+\cdots+(A_{f^{s+(a-1)p+b+1}(0)}+A_{f^{s+(a-1)p+b+2}(0)}+\cdots+A_{f^{s+ap+b}(0)})=(A_0+A_{f(0)}+\cdots+A_{f^{s+b}(0)})+aX\) が成り立ちます。 \(X\) の値及び \(A_0+A_{f(0)}+\cdots+A_{f^{s+b}(0)}\)の値を求めるにはそれぞれ\(\{ f^i(0) \}\)の最初の高々\(N\) 項を計算するだけで良く、 \(O(N)\) で求まる事から、よって十分高速にこの問題を解くことができました。

解法 2 : ダブリングを使う方法

次は、ダブリングを使う方法です。\(dp[i][j]\) を、\(A_{f^0(j)}+A_{f^1(j)}+\cdots +A_{f^{2^i-1}(j)}\) で定めます。 \(K=2^{c_1}+2^{c_2}+\cdots+2^{c_k}\) \((0\leq c_1<c_2<\cdots <c_k)\)と書けるとして、求めたいものは、\(K\) 項からなる \(A_0+A_{f(0)}+\cdots+A_{f^{K-1}(0)}\) を手前から\(2^{c_1}\)項、次の\(2^{c_2}\)項、\( \ldots\)、最後の\(2^{c_k}\)項と分けることで、 \(dp[c_1][0]+dp[c_2][s_1]+dp[c_3][s_2]+\cdots+dp[c_k][s_{k-1}]\)と求めることができます。ただし, \(s_i\)は\(\bigl(A_0+A_{f(0)}+\cdots+A_{f^{(2^{c_1}+2^{c_2}+\cdots+2^{c_{i-1}}-1)}(0)}\Bigr)\%N=(dp[c_1][0]+dp[c_2][s_1]+\cdots+dp[c_i][s_{i-1}])\%N\) で求まる値です。よって、上の\(dp[i][j]\) が求まれば良く、\(K<2^{40}\)であることから、\(0\leq i\leq 39\), \(0\leq j\leq N-1\) の範囲で求まれば良いです。

さて、\(dp[i][j]\) の値を計算する事を考えます。\(dp[0][j]=A_j\) であるので、まず初期値は簡単に計算する事が出来ます。次に、\(dp[i+1][j]\) の定義の \(2^{i+1}\) 項を前半の\(2^i\) 項と後半の\(2^i\) 項に分けることで、\(dp[i+1][j]=dp[i][j]+dp[i][(j+dp[i][j])\%N]\) と書く事が出来る事が分かります。これを用いて、\(dp[i][0],\ldots,dp[i][1],\ldots,dp[i][N-1]\) から\(dp[i+1][0],dp[i+1][1],\ldots,dp[i+1][N-1]\) を求めることができます。このとき DP を計算する部分の計算量は \(O(N\log K)\) であり、答えを計算する部分は\(O(\log K)\) であるので、十分高速です。

なお、この解法 \(2\) は解法 \(1\) に比べて、この問題については計算量が大きくなってしまっていますが、例えば \(K\) が\(Q\) 個与えられてそれぞれに対して答える問題に対しては強力です。上の解法 \(1\) を直接適用すると \(O(NQ)\) かかってしまうのに対し, こちらでは先にDP を前計算しておくことで、\(O((N+Q)\log K)\) で解くことができます。問題に応じた適切な解法選択をする必要があります。

c++による実装例 (解法1):

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

int main(void) {
	int n;
	long long k;
	long long A[200010];
	long long S[200010];
	int pre[200010];
	int s, t, p, b;
	long long a, X;
	long long ans;

	cin >> n >> k;
	for (int i = 0; i < n; i++)cin >> A[i];

	for (int i = 0; i < n; i++)pre[i] = -1;
	S[0] = 0;
	pre[0] = 0;

	for (int i = 0; i < n; i++) {
		S[i + 1] = S[i] + A[S[i] % n];
		if (pre[S[i + 1] % n] != -1) {
			s = pre[S[i + 1] % n];
			t = i + 1;
			break;
		}
		pre[S[i + 1] % n] = i + 1;
	}

	if (k <= s)ans = S[k];
	else {
		p = t - s;
		X = S[t] - S[s];
		a = (k - s - 1) / p;
		b = (k - s - 1) % p;
		ans = S[s + b + 1] + (a*X);
	}

	cout << ans << endl;
	return 0;
}

c++による実装例 (解法2):

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

int main(void) {
	int n;
	long long k;
	long long a[200010];
	long long dp[40][200010];
	long long ans;

	cin >> n >> k;
	for (int i = 0; i < n; i++)cin >> a[i];

	for (int j = 0; j < n; j++)dp[0][j] = a[j];
	for (int i = 0; i < 39; i++) {
		for (int j = 0; j < n; j++) {
			dp[i + 1][j] = dp[i][j] + dp[i][(j+dp[i][j]) % n];
		}
	}

	ans = 0;
	for (int i = 0; i < 40; i++) {
		if (k & 1) {
			ans += dp[i][ans%n];
		}
		k = (k >> 1);
	}

	cout << ans << endl;
	return 0;
}