概要

\(N\) 個のお菓子から \(1\) 個以上を選び、選んだお菓子のカロリーの総和が \(K\) の倍数になるような選び方の総数を求める問題です。

考察

お菓子の選び方を全て列挙すると \(2^N\) 通りあり、\(N\) が最大 \(2 \times 10^5\) であるため間に合いません。そこで、動的計画法（DP）を用いて効率よく数え上げることを考えます。

部分集合の総和そのものを管理すると値が大きくなりすぎますが、問題で求められているのは「総和が \(K\) の倍数かどうか」です。つまり、選んだお菓子のカロリーの総和を \(K\) で割った余りだけを管理すればよいことがわかります。

\(i\) 個目のお菓子まで選んだときの「総和を \(K\) で割った余りが \(j\) となる選び方の数」を dp[j] として管理します。 初期状態では何も選んでいないので、総和は \(0\)、すなわち余りが \(0\) の状態が \(1\) 通り（dp[0] = 1）です。

\(i\) 番目のお菓子（カロリー \(A_i\)）を処理するとき、このお菓子を「選ぶか」「選ばないか」の2択があります。 - 選ばない場合: 余りは変わりません。現在の余り \(j\) の選び方 dp[j] は、そのまま新しい next_dp[j] に引き継がれます。 - 選ぶ場合: 余り \(j\) に \(A_i \bmod K\) が加わります。新しい余りは \((j + A_i \bmod K) \bmod K\) となります。したがって、現在の dp[j] の選び方を next_dp[(j + A_i \bmod K) \bmod K] に足し合わせます。

最終的にすべてのお菓子の処理が終わった後、総和が \(K\) の倍数（余りが \(0\)）となる選び方の数は dp[0] に格納されています。ただし、問題文では「\(1\) 個以上」を選ぶ必要があるため、何も選んでいない状態（空集合、1通り）を除くために dp[0] - 1 が答えとなります。

アルゴリズム

1. DPテーブル dp をサイズ \(K\) で初期化し、dp[0] = 1 とします。
2. 各お菓子 \(i\) について、以下の処理を繰り返します。
   • 新しいDPテーブル next_dp を dp のコピーとして初期化（お菓子を選ばない場合の遷移）。
   • 各余り \(j\) (\(0 \leq j < K\)) について、next_dp[(j + A_i \bmod K) \bmod K] に dp[j] を加算する（お菓子を選ぶ場合の遷移）。
   • dp を next_dp で更新する。
3. 最後に dp[0] - 1 を出力する。ただし、答えは \(998244353\) で割った余りであるため、負の数にならないよう + 998244353 してから再度余りを取るなどの調整を行います。

計算量

• 時間計算量: \(O(N \times K)\)
• 空間計算量: \(O(K)\)

制約より \(N \times K \leq 2 \times 10^7\) であるため、\(O(N \times K)\) の計算量でも十分に実行時間内に処理が完了します。

実装のポイント

• 提供されたコードはC++で記述されていますが、Pythonの拡張子がつけられています。競技プログラミングではC++の実行速度が有利なことが多く、\(N \times K\) が \(2 \times 10^7\) に迫るこの問題ではC++の利用が適切です。

• DPテーブルの更新を行う際、毎回新しい配列 next_dp を作成していますが、これを配列のコピーで行うことで「お菓子を選ばない」場合を簡潔に表現しています。

• 余りを計算する際、idx = j + x として idx >= K なら idx -= K とすることで、% 演算（剰余演算）のオーバーヘッドを減らし、定数倍の高速化を図っています。

• 最後に dp[0] - 1 を出力しますが、dp[0] が \(0\) になることはないため、単に引き算をして負になった場合に MOD を足す処理を行っています。

  ソースコード

#include <iostream>
#include <vector>
#include <algorithm>

using namespace std;

int main() {
    ios_base::sync_with_stdio(false);
    cin.tie(NULL);
    int N, K;
    if (!(cin >> N >> K)) return 0;
    vector<int> A(N);
    for (int i = 0; i < N; ++i) {
        cin >> A[i];
    }
    
    const int MOD = 998244353;
    vector<int> dp(K, 0);
    vector<int> next_dp(K, 0);
    dp[0] = 1;
    
    for (int i = 0; i < N; ++i) {
        int x = A[i] % K;
        copy(dp.begin(), dp.end(), next_dp.begin());
        for (int j = 0; j < K; ++j) {
            int idx = j + x;
            if (idx >= K) idx -= K;
            next_dp[idx] += dp[j];
            if (next_dp[idx] >= MOD) next_dp[idx] -= MOD;
        }
        dp.swap(next_dp);
    }
    
    int ans = dp[0] - 1;
    if (ans < 0) ans += MOD;
    cout << ans << "\n";
    
    return 0;
}

---

この解説は or-glm5.2-high によって生成されました。