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 つのパターンの和になります。
- \(i+1\) 番目のお菓子を選ばない場合 余りは変化しません。\(i\) 番目までに余りが \(j\) であった状態から遷移します。
- \(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\) となります。
アルゴリズム
- 長さ \(K\) の配列
dpを用意し、dp[0] = 1、それ以外を0で初期化します。 - 各お菓子 \(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で更新します。
- 最終的な答えとして、
(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: