AtCoder Beginner Contest 262 (D)DP

1. D - I Hate Non-integer Number (DP)

https://atcoder.jp/contests/abc262/tasks/abc262_d

1.1 Problem Statement

给定长度为 $n$ 的 正整数 序列 $a_1, a_2, …, a_n$

对于每个数字 $a_i$, 可以选择, 也可以不选择

一共有 $2^n - 1$ 种选择方案(不包括全部不选择)

问: 所有的选择方案中, 有多少种方案可以使得被选择的数字的 算术平方根 为 整数 ???

1.2 分析

显而易见, 计数dp

1. 最直观的dp

   定义 $dp[i][j][k]$: 表示从序列前 $i$ 个数中选择 $j$ 个, 且和为 $k$ 的方案数

   但, 数据范围太大, 行不通

2. 改进

   修改上述定义为 $dp[i][j][k]$: 表示从序列前 $i$ 个数中选择 $j$ 个, 且 $sum \pmod j = k$ 的方案数

   但, 无法从 $dp[i][j][k]$ 转移到 $dp[i][j+1][k’]$, 转移方程不好确定, 行不通

3. 最终的dp定义

   继续修改 dp 的定义 $dp[j][k][l]$: 表示从序列前 $j$ 个数中选择 $k$ 个, 且 $sum \pmod i = l$ 的方案数

   转移方程, 对于 当前数字 $a_i$:

   • 不选择: $dp[j + 1][k][l] += dp[j][k][l]$
   • 选择: $dp[j + 1][k + 1][(l + a_j) \pmod i] += dp[j][k][l]$

1.3 code

#include <bits/stdc++.h>

using namespace std;

int main() {
    ios::sync_with_stdio(false);
    cin.tie(0);

    int n;
    cin >> n;

    vector<long long> a(n);
    for (int i = 0; i < n; i++)
    {
        cin >> a[i];
    }

    long long mod = 998244353;
    long long ans = 0;

    // 枚举每次选择的数量 [1, n]
    for (int i = 1; i <= n; i++)
    {
        // dp[j][k][l]: 从前j个数中选择k个, mod i的余数等于 l, 的方案
        vector<vector<vector<long long>>> dp(n + 1, vector<vector<long long>>(i + 1, vector<long long>(i, 0)));
        dp[0][0][0] = 1;

        for (int j = 0; j < n; j++)
        {
            for (int k = 0; k <= i; k++)
            {   
                for (int l = 0; l < i; l++)
                {
                    dp[j + 1][k][l] += dp[j][k][l];
                    dp[j + 1][k][l] %= mod;

                    if (k != i)
                    {
                        dp[j + 1][k + 1][(l + a[j]) % i] += dp[j][k][l];
                        dp[j + 1][k + 1][(l + a[j]) % i] %= mod;
                    }
                }
            }
        }
        // cout << dp[n][i][0] << endl;
        ans = (ans + dp[n][i][0] % mod) % mod;
    }

    cout << ans << endl;

    return 0;
}