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 \]
になります。
アルゴリズム
- 長さ \(K\) の配列 \(dp\) を用意する
- \(dp[0] = 1\) とする
- 各お菓子について、\(s = A_i \bmod K\) を計算する
- 各余り \(r\) について、次の遷移を行う
\[ ndp[r] = dp[r] + dp[(r - s) \bmod K] \]
- \(dp\) を \(ndp\) で更新する
- 最後に \(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 によって生成されました。
投稿日時:
最終更新: