Forzeeh4aejeech,oet a set \(S=\{ 0,1,\ldots, N-1 \}\), let us define a function \(f:S\to S\) by \(f(i)=(i+A_i)\bmod N\). Here, just as the Problem Statement, \(\mod N\) represents the remainder divided by \(N\). For any non-negative integer \(k\), the number of candies after performing \(k\) operations modulo \(N\) is \(f^k(0)\). Here, \(f^k\) is a \(k\)-time composition of \(f\), where \(f^0(x)=x\). What we want is \(A_0+A_{f(0)}+\cdots+A_{f^{K-1}(0)}\). Also, note that we can write as \(f^k(0)=\bigl( A_0+A_{f(0)}+\cdots +A_{f^{k-1}(0)}\bigr)\% N\).

Solution 1: using periodicity

Since \(f^0(0),f^1(0),\ldots,f^N(0)\) are all integers between \(0\) (inclusive) and \(n\) (exclusive), there exists \(0\leq k<k' \leq N\) such that \(f^k(0)=f^{k'}(0)\). Here, let’s take minimum \(t\) such that there exists \(0\leq s<t\) satisfying \(f^s(0)=f^t(0)\), and let \(p=t-s\). Then, for any non-negative integer \(i\), \(f^{t+i}(0)=f^i(f^t(0))=f^i(f^s(0))=f^{s+i}(0)\). That is, for any \(k\) such that \(s\leq k \), we have \(f^{k+p}(0)=f^k(0)\), so the sequence \(\{ f^i(0) \}\) has a periodicity \(p\) in the range \(s\leq i\). We use this fact to find the answer.

For \(K\leq s\), we have \(K\leq s<t\leq N\), so the transition of the number of candies can be naively simulated. Otherwise, for a non-negative integer \(a\) and an non-negative integer \(b\) less than \(p\), it can be expressed as \((K-1-s)=ap+b\). Here, for a non-negative integer \(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)} \}, \]

and by defining \(X=A_{f^s(0)}+A_{ f^{s+1}(0)}+\cdots+A_{f^{t-1}(0)}\), it holds that \(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\). The values of \(X\) and \(A_0+A_{f(0)}+\cdots+A_{f^{s+b}(0)}\) can be found from the first \(N\) terms of \(\{ f^i(0) \}\), which can be computed in a total of \(O(N)\) time, so the problem has been solved fast enough.

Solution 2: using doubling

The next solution is to use doubling. Let us define \(dp[i][j]\) by \(A_{f^0(j)}+A_{f^1(j)}+\cdots +A_{f^{2^i-1}(j)}\). When \(K=2^{c_1}+2^{c_2}+\cdots+2^{c_k}\) \((0\leq c_1<c_2<\cdots <c_k)\), the desired sum \(A_0+A_{f(0)}+\cdots+A_{f^{K-1}(0)}\) consisting of \(K\) terms can be divided into the first \(2^{c_1}\) terms, the subsequent \(2^{c_2}\), \(\ldots\), and the last \(2^{c_k}\) terms, so that the desired value can be found as \(dp[c_1][0]+dp[c_2][s_1]+dp[c_3][s_2]+\cdots+dp[c_k][s_{k-1}]\). Here, \(s_i\) is the value found as \(\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\). Therefore, all we have to find is \(dp[i][j]\). Since \(K<2^{40}\), it is sufficient to compute in the range of \(0\leq i\leq 39\) and \(0\leq j\leq N-1\).

Now, consider computing the values of \(dp[i][j]\). Since \(dp[0][j]=A_j\), the initial values can be easily found. Next, by dividing the \(2^{i+1}\) terms in the definition of \(dp[i+1][j]\) into the first \(2^i\) terms and the last \(2^i\) terms, we can write as \(dp[i+1][j]=dp[i][j]+dp[i][(j+dp[i][j])\%N]\). With this relation, we can find \(dp[i+1][0],dp[i+1][1],\ldots,dp[i+1][N-1]\) from \(dp[i][0],\ldots,dp[i][1],\ldots,dp[i][N-1]\). Here, the complexity of computing DP is \(O(N\log K)\), and that of finding the answer is \(O(\log K)\), which are fast enough.

Note that while the computational complexity of Solution \(2\) is worse than that of Solution \(1\), for the problem like answering the problem where \(Q\) number of \(K\)’s are given, this solution works better. If you apply Solution \(1\) above directly, it costs a total of \(O(NQ)\) time, but in this solution, the problem can be solved in a total of \(O((N+Q)\log K)\) by precalculating the DP. You must choose appropriate solution depending on problems.

Sample code in c++ (Solution 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;
}

Sample code in c++ (Solution 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;
}