ABC262 D - I Hate Non-integer Number

問題概要

問題文

長さ$N$の正整数列$A=(A_1,A_2,\dots,A_N)$が与えられる。
$A$から$1$個以上の要素を選ぶとき、選んだ要素すべての平均値が整数となるようなものはいくつあるか?

制約

• $1\leq N\leq 100$
• $1\leq A_i\leq 10^9$
• 入力はすべて整数

問題リンク

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

考えたこと

ある程度$N$が小さいので、$O(N^4)$ぐらいは許されることを覚えておく。
平均値と書いてあるが、要は「選んだ要素の総和」が「選んだ要素の個数」の倍数であればよい。
これは部分和問題の進化系みたいな感じなので、DPとかを使えばよさそう。
次のようなことを考えた。
dp[i][j][k]=Aの先頭i項からj個選ぶとき、選んだ要素の和をjで割った余りがkであるような個数
ただ、これだと「jで割った余り」がjが動くときに扱いづらい。
じゃあ、あらかじめ選ぶ個数sizを固定してしまえばよい。
こうすると、siz=1,2,...,Nについて、順に
「総和がsizの倍数となるようにsiz個の要素を選ぶ方法」
を求めて総和を出せばよい。これらは、次のようなDPで解決できる。

dp[i][j][k]=Aの先頭i項からj個選ぶとき、選んだ要素の和をsizで割った余りがkとなるような個数
これなら遷移が簡単。
siz=1,...,nに対してdp[n][siz][0]を求めて、それらの総和を求めればよい。



dpは状態数が$O(N^3)$で、遷移も同程度。したがって、各sizに対してdpは$O(N^3)$で埋めることができる。
さらに、sizが$O(N)$通りなので全体では$O(N^4)$の計算量で求めることができた。

実装

#include<iostream>
#include<vector>
using namespace std;
int main() {
    const int mod=998244353;

    //入力
    int n;
    cin>>n;
    vector<int> a(n);
    for(auto&aa:a)cin>>aa;
    int ans=0;

    //siz=1,2,...,nで動かす
    for(int siz=1;siz<=n;siz++){

        //dp[i][j][k]:=先頭i項からj項選ぶ。選んだ要素の和をsizで割った余りがkとなるような個数。
        vector<vector<vector<int>>> dp(n+1,vector<vector<int>>(siz+1,vector<int>(siz)));
        dp[0][0][0]=1;
        for(int i=0;i<n;i++)for(int j=0;j<=siz;j++)for(int k=0;k<siz;k++){
            //a[i]を選ばないとき
            dp[i+1][j][k]+=dp[i][j][k];
            dp[i+1][j][k]%=mod;

            //a[i]を選ぶとき
            if(j+1<=siz){
                dp[i+1][j+1][(k+a[i])%siz]+=dp[i][j][k];
                dp[i+1][j+1][(k+a[i])%siz]%=mod;
            }
        }
        ans+=dp[n][siz][0];
        ans%=mod;
    }
    //出力
    cout<<ans<<endl;
}