F - Hop Sugoroku

Problem Statement

There is a row of $N$ squares labeled $1,2,\dots,N$ and a sequence $A=(A_1,A_2,\dots,A_N)$ of length $N$.
Initially, square $1$ is painted black, the other $N-1$ squares are painted white, and a piece is placed on square $1$.

You may repeat the following operation any number of times, possibly zero:

• When the piece is on square $i$, choose a positive integer $x$ and move the piece to square $i + A_i \times x$.
  • Here, you cannot make a move with $i + A_i \times x > N$.
• Then, paint the square $i + A_i \times x$ black.

Find the number of possible sets of squares that can be painted black at the end of the operations, modulo $998244353$.

Constraints

• All input values are integers.
• $1 \le N \le 2 \times 10^5$
• $1 \le A_i \le 2 \times 10^5$

Input

The input is given from Standard Input in the following format:

$N$
$A_1$ $A_2$ $\dots$ $A_N$

Output

Print the answer as an integer.

Sample Input 1

5
1 2 3 1 1

Sample Output 1

8

There are eight possible sets of squares painted black:

• Square $1$
• Squares $1,2$
• Squares $1,2,4$
• Squares $1,2,4,5$
• Squares $1,3$
• Squares $1,4$
• Squares $1,4,5$
• Squares $1,5$

Sample Input 2

1
200000

Sample Output 2

1

Sample Input 3

40
1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1

Sample Output 3

721419738

Be sure to find the number modulo $998244353$.

解题思路

  dp 的部分还是很容易想到的。定义 $f(i)$ 表示在前 $i$ 个元素中且第 $i$ 个元素染成黑色的所有方案的数量。根据从第 $i$ 个位置出发可到达的下标来进行状态转移，因此状态转移方程就是 $f(i) \to f(j), \, j = i + k \cdot a_i, \, k \in \mathbb{Z}$。其中状态转移的计算量是 $O(\frac{n}{a_i})$ 级别，如果每个 $a_i$ 都很小，那么整个 dp 的时间复杂度就会变成 $O(n^2)$。因此只有当 $a_i$ 比较大时，这种朴素转移的计算量才会比较小。

  这时候就可以考虑用根号分治了，取 $b = 450$（大致 $\sqrt{2 \cdot 10^5}$ 级别）。当 $a_i > b$ 时直接朴素转移，那么这部分整体的时间复杂度大致是 $O(n \cdot \frac{n}{b})$。当 $n \leq b$ 时就不能直接转移了，反过来考虑 $i$ 能从哪些下标 $j, \, (j < i)$ 转移过来，即满足 $j + k \cdot a_j = i$ 的下标 $j$。如果 $a_j$ 是确定的，那么就有 $j \equiv i \pmod{a_j}$。因此用 $s[d][r]$ 来统计满足 $j \equiv r \pmod{d}$ 的 $f(j)$ 的总和，当枚举到 $i$ 时，遍历 $d \in [1, b]$，把 $s[d][i \bmod d]$ 累加到 $f(i)$ 中，表示可以从 $a_j \leq b$ 且 $j \equiv i \pmod{a_j}$ 的下标 $j$ 转移到位置 $i$。另外当 $a_i \leq b$ 时，只需让 $s[a_i][i \bmod a_i]$ 累加 $f(i)$ 即可。这部分整体的时间复杂度为 $O(n \cdot b)$。

  考虑两种情况，那么总的时间复杂度就是 $O(n \cdot \frac{n}{b} + n \cdot b)$，其中根据基本不等式 $\frac{n}{b} + b \geq 2 \sqrt{\frac{n}{b} \cdot b} = 2 \sqrt{n}$，当 $\frac{n}{b}$ 和 $b$ 相等即均取 $\sqrt{n}$ 时，等号成立，因此这里 $b$ 大概取 $450$。

  AC 代码如下，时间复杂度为 $O(n \sqrt{n})$：

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

const int N = 2e5 + 10, M = 460, mod = 998244353;

int a[N];
int f[N];
int s[M][M];

int main() {
    int n;
    scanf("%d", &n);
    for (int i = 1; i <= n; i++) {
        scanf("%d", a + i);
    }
    f[1] = 1;
    for (int i = 1; i <= n; i++) {
        for (int j = 1; j <= 450; j++) {
            f[i] = (f[i] + s[j][i % j]) % mod;
        }
        if (a[i] <= 450) {
            s[a[i]][i % a[i]] = (s[a[i]][i % a[i]] + f[i]) % mod;
        }
        else {
            for (int j = i + a[i]; j <= n; j += a[i]) {
                f[j] = (f[j] + f[i]) % mod;
            }
        }
    }
    int ret = 0;
    for (int i = 1; i <= n; i++) {
        ret = (ret + f[i]) % mod;
    }
    printf("%d", ret);
    
    return 0;
}

参考资料

  Editorial - AtCoder Beginner Contest 335 (Sponsored by Mynavi)：https://atcoder.jp/contests/abc335/editorial/9038

  暴力美学——浅谈根号分治：https://www.luogu.com.cn/blog/Amateur-threshold/pu-li-mei-xue-qian-tan-gen-hao-fen-zhi