公式

C - お菓子の詰め合わせ / Assortment of Sweets 解説 by admin

gpt-5.5-high

概要

選んだお菓子のカロリー総和が \(K\) の倍数になる、空でない部分集合の個数を求める問題です。

総和そのものではなく、「総和を \(K\) で割った余り」だけに注目して動的計画法を行います。

考察

お菓子を選ぶか選ばないかを全探索すると、選び方は \(2^N\) 通りあります。
\(N\) は最大で \(2 \times 10^5\) なので、これは到底間に合いません。

ここで重要なのは、条件が

\[ \sum_{i \in S} A_i \equiv 0 \pmod K \]

であることです。

つまり、カロリー総和の具体的な値は必要なく、\(K\) で割った余りだけ分かれば十分です。

そこで、次のような DP を考えます。

  • \(dp[r]\):これまで見たお菓子からいくつか選んだとき、カロリー総和を \(K\) で割った余りが \(r\) になる選び方の数

最初は何も選んでいない状態だけが存在するので、

\[ dp[0] = 1 \]

とします。
この \(1\) は「空集合」を表します。

お菓子を 1 個ずつ見ていき、そのお菓子を

  • 選ばない
  • 選ぶ

の 2 通りを DP に反映します。

\(i\) 番目のお菓子のカロリーを \(K\) で割った余りを

\[ s = A_i \bmod K \]

とします。

新しい余りが \(r\) になる場合は、

  • もともと余りが \(r\) で、このお菓子を選ばない
  • もともと余りが \((r - s) \bmod K\) で、このお菓子を選ぶ

のどちらかです。

したがって遷移は

\[ ndp[r] = dp[r] + dp[(r - s) \bmod K] \]

となります。

例えば \(K = 5\), \(s = 2\) のとき、新しく余り \(0\) を作るには、

  • もともと余り \(0\) で選ばない
  • もともと余り \(3\) で選ぶ

のどちらかです。
実際に \(3 + 2 \equiv 0 \pmod 5\) です。

最後に \(dp[0]\) が、総和が \(K\) で割り切れる選び方の数です。
ただし、最初に入れた空集合も含まれているため、答えは

\[ dp[0] - 1 \]

になります。

アルゴリズム

  1. 長さ \(K\) の配列 \(dp\) を用意する
  2. \(dp[0] = 1\) とする
  3. 各お菓子について、\(s = A_i \bmod K\) を計算する
  4. 各余り \(r\) について、次の遷移を行う

\[ ndp[r] = dp[r] + dp[(r - s) \bmod K] \]

  1. \(dp\)\(ndp\) で更新する
  2. 最後に \(dp[0] - 1\) を出力する

この問題では \(N \times K \leq 2 \times 10^7\) なので、\(O(NK)\) の DP が十分間に合います。

計算量

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

実装のポイント

余り \(s = 0\) の場合、現在のお菓子を選んでも余りは変わりません。

このとき各状態について、

  • 選ばない
  • 選ぶ

の 2 通りがあるだけなので、

\[ dp[r] \leftarrow 2 \times dp[r] \]

とすればよいです。

また、答えは非常に大きくなるため、計算のたびに \(998244353\) で割った余りを取ります。

最後に空集合を除くため、

\[ dp[0] - 1 \]

を出力します。
引き算で負になる可能性があるため、その場合は MOD を足して調整します。

ソースコード

#include <bits/stdc++.h>
using namespace std;

static constexpr int MOD = 998244353;

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

    int N, K;
    cin >> N >> K;

    vector<int> dp(K, 0), ndp(K, 0);
    dp[0] = 1;

    for (int i = 0; i < N; i++) {
        long long A;
        cin >> A;
        int s = A % K;

        if (s == 0) {
            for (int r = 0; r < K; r++) {
                dp[r] += dp[r];
                if (dp[r] >= MOD) dp[r] -= MOD;
            }
        } else {
            for (int r = 0; r < s; r++) {
                int v = dp[r] + dp[r - s + K];
                if (v >= MOD) v -= MOD;
                ndp[r] = v;
            }
            for (int r = s; r < K; r++) {
                int v = dp[r] + dp[r - s];
                if (v >= MOD) v -= MOD;
                ndp[r] = v;
            }
            dp.swap(ndp);
        }
    }

    int ans = dp[0] - 1;
    if (ans < 0) ans += MOD;
    cout << ans << '\n';

    return 0;
}

この解説は gpt-5.5-high によって生成されました。

投稿日時:
最終更新: