ABC362E Count Arithmetic Subsequences

传送门

Solution

设 $F_{i,len, d}$ 表示以 $a_i$ 结尾，长度为 $len$，公差为 $d$ 的等差数列数量。转移方程为

$F_{i,len,d}=\sum\limits_{j=1}^{i-1} F_{j,len-1,d}$

其中 $j$ 要满足 $a_i-a_j=d$。

因为 $d$ 可能为负数，所以 dp 数组最后一维用 map。

Code

#include <bits/stdc++.h>
using namespace std;
typedef long long ll;
const ll N = 105, mod = 998244353;

ll n, a[N], ans[N];
map<ll, ll> f[N][N];

int main()
{
    scanf("%lld", &n);
    for (int i = 1; i <= n; i++) {
        scanf("%lld", &a[i]);
        for (int j = 1; j < i; j++) {
            ll d = a[i] - a[j];
            f[i][2][d] = (f[i][2][d] + 1) % mod;
        }
        for (auto &t : f[i][2])
            ans[2] = (ans[2] + t.second) % mod;
    }
    ans[1] = n;
    for (int i = 1; i <= n; i++) {
        for (int l = 3; l <= i; l++) {
            for (int j = 1; j < i; j++) {
                ll d = a[i] - a[j];
                f[i][l][d] = (f[i][l][d] + f[j][l - 1][d]) % mod;
            }
            for (auto &t : f[i][l])
                ans[l] = (ans[l] + t.second) % mod;
        }
    }
    for (int l = 1; l <= n; l++)
        printf("%lld ", ans[l]);
    return 0;
}