AtCoder が提供しているABC（AtCoder Beginner Contest）362 E問題をC++とPythonで解いてみました。ABC362は、2024年7月13日21:00に実施されました。
AtCoder の紹介はこちらに、プログラミングの方針はこちらに記事があります。
E問題 Count Arithmetic Subsequences（Difficulty : 1255）
問題の詳細は、リンク先をご覧ください。
ABC362 E問題 Count Arithmetic Subsequences
貰うDPまたはメモ化再帰で解きます。AtCoder Problems による Difficulty は 1255 でした。
解答案
数列を後ろから処理します。「えびまラボ」様の動画が参考になりました。まず、メモ化再帰を使って解きます。次に貰うDPで解きます。
C++ プログラム例（ABC362E）
再帰関数 rec(i, j, d) を、初項 $A_i$、次項 $A_j$、長さ $d (\geqq 2)$ であるような等差数列の個数とします。返す値をメモして計算量を減らします。
- 長さが2であれば、$A_i$、$A_j$ の並びは等差数列となるため、1を返します。
- $A_j$ より後ろにある $A_k$ を全探索します。
- $A_j – A_i = A_k – A_j$ であれば、$d$ の値をひとつ減らして、再帰的に呼び出します（21、22行目）。
以下が、C++プログラムです。
#include <bits/stdc++.h>
using namespace std;
#define MOD 998244353
int n;
vector<int> a;
vector<vector<vector<int>>> memo(100, vector<vector<int>>(100, vector<int>(100, -1)));
int rec(int i, int j, int d)
{
if (d == 2) {
return 1;
}
if (memo[i][j][d] != -1) {
return memo[i][j][d];
}
int result = 0;
for (int k = j + 1; k < n; ++k) {
if (a[k] - a[j] == a[j] - a[i]) {
result = (result + rec(j, k, d - 1)) % MOD;
}
}
memo[i][j][d] = result;
return result;
}
int main()
{
cin >> n;
a.resize(n);
for (int i = 0; i < n; ++i) {
cin >> a[i];
}
vector<int> result(n + 1);
result[1] = n;
for (int i = 0; i < n; ++i) {
for (int j = i + 1; j < n; ++j) {
for (int d = 2; d <= n; ++d) {
result[d] = (result[d] + rec(i, j, d)) % MOD;
}
}
}
for (int i = 1; i <= n; ++i) {
cout << result[i] << " \n"[i == n];
}
return 0;
}貰うDPとしても書けます。DP[i][j][k] を初項 $A_i$、次項 $A_j$、長さ $d (\geqq 2)$ であるような等差数列の個数とします。$i$ の値を後ろから処理します（16行目）。
#include <bits/stdc++.h>
using namespace std;
#define MOD 998244353
int main()
{
int n;
cin >> n;
vector<int> a(n);
for (int i = 0; i < n; ++i) {
cin >> a[i];
}
vector<vector<vector<int>>> dp(n, vector<vector<int>>(n, vector<int>(n + 1, 0)));
for (int i = n - 2; i >= 0; --i) {
for (int j = i + 1; j < n; ++j) {
dp[i][j][2] += 1;
dp[i][j][2] %= MOD;
for (int d = 2; d < n; ++d) {
for (int k = j + 1; k < n; ++k) {
if (a[j] - a[i] == a[k] - a[j]) {
dp[i][j][d + 1] += dp[j][k][d];
dp[i][j][d + 1] %= MOD;
}
}
}
}
}
vector<int> result(n + 1);
result[1] = n;
for (int i = 0; i < n; ++i) {
for (int j = 0; j < n; ++j) {
for (int d = 2; d <= n; ++d) {
result[d] += dp[i][j][d];
result[d] %= MOD;
}
}
}
for (int i = 1; i <= n; ++i) {
cout << result[i] << " \n"[i == n];
}
return 0;
}AC（Accepted＝正しいプログラム）と判定されました。
Python プログラム例（ABC362E）
C++のDP版を移植しました。
以下のプログラムは、「Python (CPython 3.11.4)」では、TLE判定となり、「Python (PyPy 3.10-v7.3.12)」で、AC判定となりました。
"""AtCoder Beginner Contest 362 E"""
MOD = 998244353
n = int(input())
a = list(map(int, input().split()))
dp = [[[0 for _ in range(n + 1)] for _ in range(n)] for _ in range(n)]
for i in range(n - 2, -1, -1):
for j in range(i + 1, n):
dp[i][j][2] += 1
dp[i][j][2] %= MOD
for d in range(2, n):
for k in range(j + 1, n):
if a[j] - a[i] == a[k] - a[j]:
dp[i][j][d + 1] += dp[j][k][d]
dp[i][j][d + 1] %= MOD
result = [0] * (n + 1)
result[1] = n
for i in range(n):
for j in range(n):
for d in range(2, n + 1):
result[d] += dp[i][j][d]
result[d] %= MOD
print(*result[1:])こちらも「AC」と判定されました。
最後に
メモ化再帰と貰うDPは、どちらも書けるようになりたいと考えています。よい練習になりました。
引き続き ABC の問題を紹介していきます。
