概要

\(N\) 個のお菓子から \(1\) 個以上を選び、選んだカロリーの総和が \(K\) の倍数になるような選び方の数を求める問題です。答えを \(998244353\) で割った余りで出力します。

考察

素朴なアプローチの問題点

すべての部分集合を列挙すると \(2^N\) 通りあり、\(N \leq 2 \times 10^5\) では到底間に合いません。指数時間のアプローチは現実的ではありません。

重要な気づき

ここで注目すべきは、「カロリーの総和が \(K\) で割り切れるか」を判定するには、総和を \(K\) で割った余りだけを覚えておけばよいという点です。総和そのものは巨大になりますが、余りは \(0\) から \(K-1\) の \(K\) 通りしかありません。

そこで、「いくつかのお菓子を選んだとき、その総和を \(K\) で割った余りが \(r\) になるような選び方の数」 を状態として管理する動的計画法（DP）を考えます。

具体的には、各お菓子を「選ぶ／選ばない」と順に決めていき、それぞれの余りの状態がどれだけの選び方で実現できるかを更新していきます。

アルゴリズム

配列 dp[r] を「これまで処理したお菓子の中から選んだとき、総和を \(K\) で割った余りが \(r\) になる選び方の数」と定義します。

初期状態

まだ何も選んでいない状態（空集合）では総和は \(0\) なので、dp[0] = 1、それ以外は \(0\) とします。

遷移

お菓子 \(i\)（カロリーを \(K\) で割った余りを \(a\) とする）を処理するとき、各選び方について以下の \(2\) 通りが生じます。

• このお菓子を選ばない: 余りは変わらない
• このお菓子を選ぶ: 余りが \(a\) だけ増える（\(K\) で割った余りで考える）

したがって新しい dp は、古い dp に対して以下のように更新されます。

\[ \text{new\_dp}[(r + a) \bmod K] \mathrel{+}= \text{dp}[r] \]

これを「選ばない場合（そのまま）」と合わせて整理すると、余りが \(r\) になる新しい選び方は次の \(2\) つの寄与の和になります。

\[ \text{new\_dp}[r] = \text{dp}[r] + \text{dp}[(r - a) \bmod K] \]

• 第 \(1\) 項: お菓子 \(i\) を選ばなかった分（もとから余り \(r\)）
• 第 \(2\) 項: お菓子 \(i\) を選んだ分（選ぶ前の余りが \((r-a)\) だったもの）

\(a = 0\) の特別な場合

\(a = 0\)（カロリーが \(K\) の倍数）のときは、選んでも選ばなくても余りは変わりません。よってすべての状態が \(2\) 倍になります。

最終的な答え

すべてのお菓子を処理した後、dp[0] が「総和が \(K\) で割り切れる選び方の数」になります。ただしこれには空集合（何も選ばない場合）も含まれているため、\(1\) を引いた dp[0] - 1 が答えです。

コードでの工夫

第 \(2\) 項の \(\text{dp}[(r-a) \bmod K]\) は、配列を \(s = K - a\) だけ巡回シフト（ローテーション）したものに相当します。コードでは

s = K - a
shifted = dp[s:] + dp[:s]

として、シフト後の配列 shifted[r] が dp[(r-a) mod K] を表すようにし、dp[r] + shifted[r] で一括更新しています。

計算量

• 時間計算量: \(O(N \times K)\)
• 空間計算量: \(O(K)\)

お菓子 \(1\) 個あたりの更新が \(O(K)\) で、それを \(N\) 個分行うため \(O(N \times K)\) となります。制約 \(N \times K \leq 2 \times 10^7\) により、この計算量で十分高速に動作します。

実装のポイント

• カロリー \(A_i\) は最大 \(10^9\) と大きいですが、DPで必要なのは \(A_i \bmod K\) のみなので、最初に余りを取ってしまえば問題ありません。

• 各更新でこまめに \(\bmod\ 998244353\) を取り、オーバーフローや値の肥大化を防ぎます（Python では多倍長整数のため計算が遅くならないようにする目的もあります）。

• 最後に空集合の分を忘れずに \(1\) を引きます。(dp[0] - 1) % MOD とすることで、万一 dp[0] が \(0\) でも負にならないよう剰余を取っています。

  ソースコード

import sys

def main():
    data = sys.stdin.buffer.read().split()
    N = int(data[0]); K = int(data[1])
    A = data[2:2+N]
    MOD = 998244353
    dp = [0] * K
    dp[0] = 1
    for x in A:
        a = int(x) % K
        if a == 0:
            dp = [(v * 2) % MOD for v in dp]
        else:
            s = K - a
            shifted = dp[s:] + dp[:s]
            dp = [(p + q) % MOD for p, q in zip(dp, shifted)]
    print((dp[0] - 1) % MOD)

main()

---

この解説は claude4.8opus-high によって生成されました。