公式

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

claude4.8opus-high

概要

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

考察

まず素朴に考えると、\(N\) 個のお菓子それぞれを「選ぶ・選ばない」の \(2\) 通りなので、部分集合は全部で \(2^N\) 通りあります。これを全列挙して総和が \(K\) の倍数かどうか確認する方法は、\(N\) が最大 \(2 \times 10^5\) のため、到底間に合いません(指数時間)。

ここで重要な気づきは、「総和が \(K\) で割り切れるか」を判定するには、総和そのものではなく、総和を \(K\) で割った余りだけが分かればよいという点です。

たとえば \(K = 3\) のとき、カロリー \(7\) のお菓子と \(4\) のお菓子は、\(K\) で割った余りがそれぞれ \(1\) で同じです。総和が \(K\) の倍数になるかどうかを考える上では、各お菓子のカロリーを「\(K\) で割った余り」に置き換えてしまっても答えは変わりません。

そこで、選んだお菓子の総和を \(K\) で割った余り をキーとして、選び方の数を管理していくことを考えます。余りは \(0, 1, \ldots, K-1\)\(K\) 通りしかないので、これを状態として動的計画法(DP)が組めます。

アルゴリズム

dp[j] を「ここまで処理したお菓子の中から部分集合を選び、その総和を \(K\) で割った余りが \(j\) になる選び方の数」と定義します。

  • 初期状態:何も選ばない(空集合)ときの余りは \(0\) なので、dp[0] = 1、それ以外は \(0\) とします。

次に、お菓子を \(1\) つずつ処理します。\(i\) 番目のお菓子のカロリーの余りを \(r = A_i \bmod K\) とすると、各お菓子について「選ばない」か「選ぶ」かの \(2\) 通りがあります。

新しい状態 ndp[j] は、

  • このお菓子を選ばない場合:これまでの余りが \(j\) だったものがそのまま余り \(j\) になる → dp[j]
  • このお菓子を選ぶ場合:これまでの余りが \((j - r) \bmod K\) だったものに \(r\) を足すと余りが \(j\) になる → dp[(j - r) \bmod K]

の和になります。つまり

\[ \text{ndp}[j] = \text{dp}[j] + \text{dp}[(j - r) \bmod K] \]

これを全お菓子について繰り返します。

最後に dp[0] は「総和が \(K\) の倍数になる選び方の数」ですが、ここには空集合(何も選ばない)も含まれています。問題は「\(1\) 個以上選ぶ」ことを要求しているので、空集合の分 \(1\) を引いた dp[0] - 1 が答えになります。

具体例

\(N = 3\), \(K = 3\), \(A = (1, 2, 3)\) の場合を考えます。余りはそれぞれ \(1, 2, 0\) です。

  • 初期:dp = [1, 0, 0]
  • お菓子1(\(r=1\))処理後:dp = [1, 1, 0](空 / {1})
  • お菓子2(\(r=2\))処理後:dp = [1, 1, 1] … 余り0は{}と{1,2}、余り1は{1}と{2}…のように積み上がる
  • お菓子3(\(r=0\))処理後:dp = [2, 2, 2]

余り \(0\) になるのは dp[0] = 2(空集合と \(\{1,2\}\))。空集合を除くと答えは \(1\) です。実際、総和が \(3\) の倍数になるのは \(\{1,2\}\)(総和 \(3\))のみで、正しいです。

計算量

  • 時間計算量: \(O(N \times K)\)
    • お菓子ごとに \(K\) 個の余りについて更新するため。制約より \(N \times K \leq 2 \times 10^7\) なので十分高速です。
  • 空間計算量: \(O(K)\)
    • DP 配列を \(2\) 本(dpndp)持つだけで済みます。

実装のポイント

  • 余りのマイナスに注意\((j - r) \bmod K\) を計算する際、\(j - r\) が負になることがあります。C++ の % 演算子は負の値に対して負を返すため、if(prev < 0) prev += K; のように補正が必要です。

  • MOD を取りながら計算:答えは非常に大きくなるため、加算のたびに \(998244353\) で剰余を取ります。コードでは加算後に if(ndp[j] >= MOD) ndp[j] -= MOD; として、剰余演算より高速に処理しています。

  • 配列の使い回しswap(dp, ndp) でポインタを入れ替えることで、毎回新しい配列を確保するコストを避けられます。ndp を毎回上書きするので初期化漏れの心配もありません。

  • 空集合の除外:最後に dp[0] - 1 とすることを忘れないようにしましょう。引き算で負にならないよう (dp[0] - 1 + MOD) % MOD としています。

    ソースコード

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

const long long MOD = 998244353;

int main(){
    int N, K;
    scanf("%d %d", &N, &K);
    
    vector<long long> dp(K, 0), ndp(K, 0);
    dp[0] = 1; // empty subset
    
    for(int i = 0; i < N; i++){
        long long a;
        scanf("%lld", &a);
        int r = (int)(a % K);
        // ndp[j] = dp[j] + dp[(j - r) mod K]
        for(int j = 0; j < K; j++){
            int prev = j - r;
            if(prev < 0) prev += K;
            ndp[j] = dp[j] + dp[prev];
            if(ndp[j] >= MOD) ndp[j] -= MOD;
        }
        swap(dp, ndp);
    }
    
    long long ans = (dp[0] - 1 + MOD) % MOD;
    printf("%lld\n", ans);
    
    return 0;
}

この解説は claude4.8opus-high によって生成されました。

投稿日時:
最終更新: