C - お菓子の詰め合わせ / Assortment of Sweets Editorial by admin Editorial for printing
gpt-5.5-high概要
各お菓子を選ぶか選ばないかを考え、選んだお菓子のカロリー総和が \(K\) の倍数になる空でない部分集合の個数を求めます。
重要なのは、総和そのものではなく「\(K\) で割った余り」だけを管理すればよい、という点です。
考察
部分集合は全部で \(2^N\) 個あります。
\(N\) は最大で \(2 \times 10^5\) なので、すべての部分集合を列挙する方法は到底間に合いません。
そこで、各お菓子のカロリー \(A_i\) について、実際の値ではなく
\(A_i \bmod K\)
だけに注目します。
なぜなら、ある部分集合の総和が \(K\) で割り切れるかどうかは、その総和の余りが \(0\) かどうかだけで決まるからです。
例えば \(K = 5\) のとき、
- カロリー \(12\) は余り \(2\)
- カロリー \(7\) も余り \(2\)
なので、どちらも「総和の余りを \(2\) 増やすお菓子」として扱えます。
ここで、以下のような DP を考えます。
\(dp[j] =\) これまで見たお菓子の中からいくつか選んで、総和の余りが \(j\) になる選び方の数
初期状態では、何も選ばない場合があるので
\(dp[0] = 1\)
です。
余りが \(r\) のお菓子を 1 個追加で考えると、各部分集合について
- そのお菓子を選ばない
- そのお菓子を選ぶ
の 2 通りがあります。
最終的な余りを \(j\) とすると、
- 選ばない場合: もともと余り \(j\)
- 選ぶ場合: もともと余り \((j-r) \bmod K\)
だった必要があります。
したがって更新式は
\(next[j] = dp[j] + dp[(j-r) \bmod K]\)
となります。
また、余りが \(0\) のお菓子は特別です。
余り \(0\) のお菓子を選んでも、総和の余りは変わりません。
そのため、余り \(0\) のお菓子が \(c\) 個あるなら、それらは独立に選ぶ・選ばないを決められるので、全体の答えを \(2^c\) 倍すればよいです。
最後に、DP には「何も選ばない」場合も含まれているため、空集合の 1 通りを引きます。
アルゴリズム
- 各 \(A_i\) について \(A_i \bmod K\) を計算し、余りごとの個数を数える。
- \(dp[0] = 1\)、それ以外を \(0\) とする。
- 余りが \(1\) 以上の各お菓子について、DP を更新する。
- 余りが \(r\) のお菓子を 1 個処理するとき、
$\( next[j] = dp[j] + dp[(j-r) \bmod K] \)$
とする。 4. 余り \(0\) のお菓子の個数を \(cnt[0]\) とすると、余りを変えずに自由に選べるので、答えは
$\( dp[0] \times 2^{cnt[0]} - 1 \)$
となる。 5. 答えを \(998244353\) で割った余りで出力する。
具体例として、\(K=5\)、余りが \(2\) のお菓子を処理する場合を考えます。
このお菓子を選んだ結果、余りが \(3\) になるためには、もともとの余りは
\[ (3 - 2) \bmod 5 = 1 \]
である必要があります。
そのため、
\[ next[3] = dp[3] + dp[1] \]
となります。
計算量
- 時間計算量: \(O(NK)\)
- 空間計算量: \(O(K)\)
制約で \(N \times K \leq 2 \times 10^7\) と保証されているため、この DP は十分高速に実行できます。
実装のポイント
- DP 配列は長さ \(K\) で持ちます。
- Python では
old[j - r]と書くと、\(j-r\) が負の場合に末尾から参照されます。これはちょうど \(\bmod K\) の扱いと一致します。 - 余り \(0\) のお菓子は DP で処理せず、最後に
pow(2, cnt[0], MOD)を掛けます。 - 空集合を除く必要があるため、最後に \(1\) を引きます。
- \(K=1\) の場合は、すべての部分集合の総和が必ず \(1\) で割り切れるので、答えは
$\( 2^N - 1 \)$
です。
ソースコード
import sys
MOD = 998244353
def main():
data = list(map(int, sys.stdin.buffer.read().split()))
N, K = data[0], data[1]
if K == 1:
print((pow(2, N, MOD) - 1) % MOD)
return
cnt = [0] * K
residues = []
for a in data[2:]:
r = a % K
if cnt[r] == 0:
residues.append(r)
cnt[r] += 1
dp = [0] * K
dp[0] = 1
m = MOD
rk = range(K)
for r in residues:
if r == 0:
continue
c = cnt[r]
for _ in range(c):
old = dp
dp = [(v if (v := old[j] + old[j - r]) < m else v - m) for j in rk]
ans = (dp[0] * pow(2, cnt[0], m) - 1) % m
print(ans)
if __name__ == "__main__":
main()
この解説は gpt-5.5-high によって生成されました。
posted:
last update: