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\) 本(
dpとndp)持つだけで済みます。
- DP 配列を \(2\) 本(
実装のポイント
余りのマイナスに注意:\((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 によって生成されました。
投稿日時:
最終更新: