Official

C - お菓子の詰め合わせ / Assortment of Sweets Editorial by admin Editorial for printing

gemini-3.5-flash-high

概要

\(N\) 個のお菓子からいくつか(1個以上)を選び、そのカロリーの総和が \(K\) の倍数となるような選び方の数を \(998244353\) で割った余りで求める問題です。

考察

1. 素朴なアプローチとその限界

お菓子の選び方は、各お菓子を「選ぶ」か「選ばない」かの \(2\) 択であるため、全体で \(2^N\) 通り存在します。 \(N \le 2 \times 10^5\) という制約では、\(2^N\) はあまりにも膨大な数になり、全探索(すべての組み合わせを試すこと)では実行時間制限(TLE)になってしまいます。

2. 「余り」に着目した動的計画法 (DP)

「総和が \(K\) の倍数」という条件は、「総和を \(K\) で割った余りが \(0\)」と言い換えることができます。 お菓子を 1 つずつ順番に選ぶかどうか決めていくとき、それまでに選んだお菓子の「具体的な総和」をすべて覚えておく必要はありません。「総和を \(K\) で割った余り」さえ管理しておけば、次のお菓子を加えたときの余りを正しく計算できます。

このように、状態を「余り」にまとめることで、動的計画法 (DP) を適用することができます。

3. DP の設計

以下のような DP テーブルを定義します。

  • \(dp[i][j]\) : 最初から \(i\) 番目までのお菓子の中からいくつか選んだとき、カロリーの総和を \(K\) で割った余りが \(j\) になるような選び方の数

余りは \(0\) から \(K-1\) までの \(K\) 通りです。 \(i+1\) 番目のお菓子(カロリー \(A_{i+1}\)、余り \(rem = A_{i+1} \bmod K\))を考えるとき、遷移は以下の 2 つのパターンの和になります。

  1. \(i+1\) 番目のお菓子を選ばない場合 余りは変化しません。\(i\) 番目までに余りが \(j\) であった状態から遷移します。
  2. \(i+1\) 番目のお菓子を選ぶ場合 選ぶ前の余りを \(prev\) とすると、 \(prev + rem \equiv j \pmod K\) となる必要があります。 これを逆算すると、 \(prev \equiv j - rem \pmod K\) となります。

したがって、遷移式は以下のようになります。 $\(dp[i+1][j] = (dp[i][j] + dp[i][(j - rem + K) \bmod K]) \bmod 998244353\)$

初期状態は、お菓子を 1 つも選んでいない状態(総和が \(0\)、すなわち余り \(0\))が \(1\) 通り存在するので、以下のように設定します。 * \(dp[0][0] = 1\) * \(dp[0][j] = 0 \quad (1 \le j < K)\)

4. 空集合の除外

求めたいのは「1個以上選ぶ」方法の数です。 DP の計算には「お菓子を 1 つも選ばない(空集合)」という \(1\) 通り(総和 \(0\)、余り \(0\))が含まれてしまっています。 そのため、最終的な答えは \(dp[N][0] - 1\) となります。

アルゴリズム

  1. 長さ \(K\) の配列 dp を用意し、dp[0] = 1、それ以外を 0 で初期化します。
  2. 各お菓子 \(A_i\) (\(1 \le i \le N\)) について順番に以下を行います。
    • \(A_i\) を \(K\) で割った余り \(rem = A_i \bmod K\) を求めます。
    • 次の状態を表す配列 next_dp を用意します。
    • 各 \(j\) (\(0 \le j < K\)) について、次の遷移を行います。 next_dp[j] = (dp[j] + dp[(j - rem + K) % K]) % MOD
    • dp を next_dp で更新します。
  3. 最終的な答えとして、(dp[0] - 1 + MOD) % MOD を出力します。

計算量

  • 時間計算量: \(O(NK)\) お菓子 \(N\) 個に対して、それぞれ \(K\) 回の遷移を行います。制約より \(N \times K \le 2 \times 10^7\) であるため、実行時間制限に余裕で間に合います。
  • 空間計算量: \(O(K)\) 2次元配列 \(dp[N+1][K]\) をそのまま保持するとメモリを多く消費しますが、遷移には直前の行の情報(\(i\) 番目の状態)しか使わないため、1次元配列を使い回すことで空間計算量を \(O(K)\) に抑えることができます。

実装のポイント

  • メモリの節約: C++の実装コードにある通り、サイズ \(K\) の配列を2つ(dp と next_dp)用意して交互に更新していくことで、メモリ使用量を大幅に削減しています。

  • 高速化(剰余演算の回避): 遷移の際、% MOD の演算は比較的低速です。足し算の結果が MOD を超えた場合に MOD を引く処理(if (val_next >= MOD) val_next -= MOD;)に置き換えることで、実行時間を大幅に短縮できます。

  • 負の数の余り: 最後に dp[0] - 1 を計算する際、値が \(0\) の場合は結果が \(-1\) になってしまいます。余りの計算で負の数にならないよう、MOD を足してから MOD で割った余りを取るように注意しましょう。

    ソースコード

#include <iostream>
#include <vector>

using namespace std;

int main() {
    // 高速な入出力
    ios_base::sync_with_stdio(false);
    cin.tie(NULL);

    int N, K;
    if (!(cin >> N >> K)) return 0;

    const int MOD = 998244353;
    vector<int> dp(K, 0);
    dp[0] = 1; // 空集合(総和0)の1通り

    vector<int> next_dp(K);

    for (int i = 0; i < N; ++i) {
        long long val;
        cin >> val;
        int rem = val % K;
        for (int j = 0; j < K; ++j) {
            int prev = j - rem;
            if (prev < 0) prev += K;
            int val_next = dp[j] + dp[prev];
            if (val_next >= MOD) val_next -= MOD;
            next_dp[j] = val_next;
        }
        dp = next_dp;
    }

    // 空でない部分集合を求めるため、何も選ばない1通りを引く
    int ans = dp[0] - 1;
    if (ans < 0) ans += MOD;

    cout << ans << "\n";

    return 0;
}

この解説は gemini-3.5-flash-high によって生成されました。

posted:
last update: