問題概要
左から右に一列に並んだ $n$ 個のマスがあり，$1, 2, \dots, n$ で番号付けられている．また，長さ $n$ の数列 $A = \langle A_1, A_2, \dots, A_n \rangle$ がある．
初期状態では，マス $1$ は黒く塗られていて，その他のマスは白く塗られている．またマス $1$ に駒が $1$ つ置かれている．
ここで，次の一連の操作を $0$ 回以上繰り返すことができる．
- 駒がマス $i$ にあるとする．正整数 $x \in \mathbb Z^+$ を選び，駒をマス $i + xA_i$ に移動させる．
- 盤面をはみ出すような $x$ を選ぶことはできない．
- マス $i + xA_i$ を黒く塗る．
操作を終えた時点で黒く塗られているマスの集合としてあり得るものの個数を $r$ として，$r \bmod 998{,}244{,}353$ を求めよ．
制約
- $1 \leq n \leq 2 \times 10^5$
- $1 \leq A_i \leq 2 \times 10^5$
解法
まず重要な言い換え*1があります．「黒く塗られたマスの集合」というのは「マス $1$ から出発して通ったマスの集合」ということで，「移動の仕方」と一対一で対応します．なので，黒く塗られたマスの内で一番右にあるものがマス $i$ であるような塗られ方は，マス $1$ からマス $i$ へ移動する方法の総数と一致します．よって，全てのマスについてこの値を求めることができれば，その総和をとることで問題を解くことができます．
DP に慣れた人なら思いつくであろう DP として，$$\mathit{ dp }[i] := \text{ マス $i$ を終点とする移動方法の総数 }$$ という風に状態をとる DP があります．この DP の状態遷移は，
- $i$ の昇順に，$\mathit{ dp }[i]$ を有効な各 $\mathit{ dp }[ i + xA_i ]$ へ足し込む．
です．
この DP は欲しい値を正しく計算できますが，時間計算量の面で問題があります．遷移の総数が最大になるケースを考えてみると，全ての $i$ について $A_i = 1$ のときに遷移の数が最大になり，その数は $\Theta( n^2 )$ 個あります．制約からこれでは TLE するので，より速い方法を考える必要があります．
ところで，マス $i$ からの遷移の数は $O\left( \frac n { A_i } \right)$ 個であると言えるので，各 $A_i$ がある程度大きければ，上記の DP が間に合います．ということは，小さい $A_i$ に関する処理だけを上手く高速化できればよいということになります．
小さい $A_i$ に対する処理
高速化のために，状態遷移についてもう少し深く考えてみます．$\mathit{ dp }[i]$ から $\mathit{ dp }[j]$ への更新が発生したとき，題意から $j - i$ は $A_i$ の倍数です．「差が何々の倍数」というのは，合同式の性質から「両者が何々を法として合同である」と言い換えられます．つまり $i \equiv j \pmod { A_i } $ です．このことから，$j$ の方から見ると $\mathit{ dp }[j]$ への更新が発生したときの $i$ とは，ある正整数 $k \in \mathbb Z^+$ であって，$i \bmod k = j \bmod k$ となるようなものが存在するような $i$ です．ある $j$ を見ているときに $k$ の値を決め打ちしてあげれば対応する $i$ が自動的に決まるので，あり得る $k$ 全てについて剰余ごとに個数をキャッシュしておくことで，いわゆる「貰う DP」をするときの更新元になる要素の和を前計算できます．具体的には，$\mathit{ dp }[i]$ から「配る DP」をする代わりに，「DP 配列に未反映の値であって $k$ での剰余が $l$ なもの」を $c_{ k, l }$ として溜めておいて，更新先に配る代わりに $$c_{ A_i, i \bmod { A_i } } \leftarrow c_{ A_i, i \bmod { A_i } } + \mathit{ dp }[i]$$ をして，$\mathit{ dp }[j]$ を見るときに $$\mathit{ dp }[j] \leftarrow \mathit{ dp }[j] + c_{ k, j \bmod k } \quad \text{(for all valid $k$)}$$ をして反映します．
「小さい」$A_i$ とは？
上記のキャッシュ戦略を実現するにあたり，各 $k$ について $k$ での剰余ごとに値を保存する必要があるので，$k$ の最大値を $m $ とすると $\Theta( m^2 )$ 空間が必要になり，DP 配列への反映時に $\Theta( m )$ 時間がかかります．一方で $m $ が小さすぎると，愚直に配るときの処理が間に合わなくなります．ではどのような値を閾値にすればよいか考えてみると，各 $i$ について
- 愚直に配る処理は $O( \frac n { \min\{ \{ A_j \mid m \leq A_j \} \} } ) = O( \frac n m )$ 時間かかる
- キャッシュ戦略は $\Theta( m^2 )$ 空間と $\Theta( m )$ 時間かかる
ということなので，この両者のバランスを取ります．やや天啓かもですが，$m = \sqrt n$ とするとどちらも $O( \sqrt n )$ 時間，$\Theta( n )$ 空間となっていい感じになり，全体では $O( n \sqrt n )$ 時間となって間に合います．
コード
AtCoder Library (ACL) を利用しています．
#include <atcoder/modint> using namespace atcoder; using MINT = modint998244353; int main() { IN( int, N ); VI A( N ); cin >> A; const int sqN = sqrt( N ); static MINT dp[ 2 << 18 ]; dp[0] = 1; auto cache = make_vector< MINT >( sqN + 1, sqN, 0 ); REP( i, N ) { // restore cache REP( j, 1, sqN ) { dp[i] += cache[j][ i % j ]; } if ( sqN <= A[i] ) { for ( int x = 1; i + x * A[i] < N; ++x ) { dp[ i + x * A[i] ] += dp[i]; } } else { cache[ A[i] ][ i % A[i] ] += dp[i]; } } cout << accumulate( ALL( dp ), MINT( 0 ) ).val() << endl; return 0; }
*1:と言うほど非自明ではないという意見もありそうですが．
