Code : Hop Sugoroku
#include <atcoder/modint>
#include <bits/stdc++.h>
using namespace std;
using namespace atcoder;
using mint = modint998244353;
int solve(vector<long long> &a) {
int n = a.size();
int rootn = sqrt(n);
vector<mint> dp(n);
vector<vector<mint>> delta(n, vector<mint>(rootn + 1, 0));
dp[0] = 1;
for (int i = 0; i < n; i++) {
for (int j = 1; j <= rootn; j++) {
dp[i] += delta[i][j];
}
if (a[i] > rootn) {
for (int j = i + a[i]; j < n; j += a[i]) {
dp[j] += dp[i];
}
} else if (i + a[i] < n) {
delta[i + a[i]][a[i]] += dp[i];
}
for (int j = 1; j <= rootn; j++) {
if (i + j < n) {
delta[i + j][j] += delta[i][j];
}
}
}
mint ans = 0;
for (int i = 0; i < n; i++) {
ans += dp[i];
}
return ans.val();
}
int main() {
int n;
cin >> n;
vector<long long> a(n);
for (int i = 0; i < n; i++) {
cin >> a[i];
}
cout << solve(a) << endl;
return 0;
}
#include <atcoder/modint>
#include <bits/stdc++.h>
using namespace std;
using namespace atcoder;
using mint = modint998244353;
int solve(vector<long long> &a) {
int n = a.size();
vector<mint> dp(n);
dp[0] = 1;
for (int i = 0; i < n; i++) {
for (int j = i + a[i]; j < n; j += a[i]) {
dp[j] += dp[i];
}
}
mint ans = 0;
for (int i = 0; i < n; i++) {
ans += dp[i];
}
return ans.val();
}
int main() {
int n;
cin >> n;
vector<long long> a(n);
for (int i = 0; i < n; i++) {
cin >> a[i];
}
cout << solve(a) << endl;
return 0;
}
#include <atcoder/modint>
#include <bits/stdc++.h>
using namespace std;
using namespace atcoder;
using mint = modint998244353;
int solve(vector<long long> &a) {
int n = a.size();
vector<mint> dp(n);
vector<vector<mint>> delta(n, vector<mint>(n + 1, 0));
dp[0] = 1;
for (int i = 0; i < n; i++) {
for (int j = 1; j <= n; j++) {
dp[i] += delta[i][j];
}
if (i + a[i] < n) {
delta[i + a[i]][a[i]] += dp[i];
}
for (int j = 1; j <= n; j++) {
if (i + j < n) {
delta[i + j][j] += delta[i][j];
}
}
}
mint ans = 0;
for (int i = 0; i < n; i++) {
ans += dp[i];
}
return ans.val();
}
int main() {
int n;
cin >> n;
vector<long long> a(n);
for (int i = 0; i < n; i++) {
cin >> a[i];
}
cout << solve(a) << endl;
return 0;
}
This submission is also correct. But it’s non trivial to prove that the time complexity is O(n sqrt(n))
. See this editorial for a proof.
#include <atcoder/modint>
#include <bits/stdc++.h>
using namespace std;
using namespace atcoder;
using mint = modint998244353;
int solve(vector<long long> &a) {
int n = a.size();
// dp[i] is the number of black subsets of the first i elements such that
// the i-th element is necessarily painted black.
vector<mint> dp(n), offload(n);
dp[0] = 1;
for (int i = 0; i < n; i++) {
for (int j = i + a[i]; j < n; j += a[i]) {
dp[j] += dp[i] + offload[i];
if (a[j] == a[i]) {
offload[j] += offload[i] + dp[i];
break;
}
}
}
mint ans = 0;
for (int i = 0; i < n; i++) {
ans += dp[i];
}
return ans.val();
}
int main() {
int n;
cin >> n;
vector<long long> a(n);
for (int i = 0; i < n; i++) {
cin >> a[i];
}
cout << solve(a) << endl;
return 0;
}