公式

C - お菓子の詰め合わせ / Assortment of Sweets 解説 by admin 印刷ページ

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 通りを引きます。

アルゴリズム

  1. 各 \(A_i\) について \(A_i \bmod K\) を計算し、余りごとの個数を数える。
  2. \(dp[0] = 1\)、それ以外を \(0\) とする。
  3. 余りが \(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 によって生成されました。

投稿日時:
最終更新: