概要

\(N\) 個のお菓子からいくつか（1個以上）を選び、そのカロリーの総和が \(K\) の倍数になるような選び方の数を求める問題です。 動的計画法（DP）を用いることで、制約の \(N \times K \le 2 \times 10^7\) という性質を活かして高速に解くことができます。

考察

1. カロリーの余りに注目する

選んだお菓子の総和が \(K\) で割り切れるかどうかを判定するためには、各お菓子のカロリー \(A_i\) をそのまま扱う必要はありません。カロリーを \(K\) で割った余り \(A_i \bmod K\) のみに関心があります。したがって、あらかじめすべてのお菓子のカロリーを \(K\) で割った余りに変換して考えます。

2. 素朴なDP（動的計画法）

余りに注目すると、以下のようなDPを考えることができます。 - \(dp[i][j]\) ：最初から \(i\) 個のお菓子の中からいくつか選んで、カロリーの総和を \(K\) で割った余りが \(j\) になるような選び方の数

\(i\) 番目のお菓子（余り \(a = A_i \bmod K\)）を選ぶか選ばないかで遷移を考えます。 - 選ばない場合: \(dp[i-1][j]\) 通り - 選ぶ場合: \(dp[i-1][(j - a) \bmod K]\) 通り

よって、遷移式は次のようになります。 $\(dp[i][j] = dp[i-1][j] + dp[i-1][(j - a) \bmod K]\)$

このDPを愚直に行うと、状態数が \(N \times K\)、遷移が \(O(1)\) なので、全体で \(O(NK)\) の計算量になります。制約より \(N \times K \le 2 \times 10^7\) であるため、この計算量で実行時間制限に間に合わせることができます。

3. さらに高速化するための工夫

Pythonなどの言語では、ループの回数をできるだけ減らすことが実行時間制限（TLE）を避けるために重要です。そこで、以下の2つの工夫を行います。

① \(K = 1\) の例外処理

\(K = 1\) のとき、どのようなお菓子の組み合わせを選んでも総和は必ず \(1\) の倍数（＝任意の整数）になります。 したがって、空でない部分集合の個数は単純に \(2^N - 1\) 通りとなり、DPをせずに \(O(\log N)\) で答えを求められます。

② 余りが \(0\) のお菓子を別で処理する

カロリーが最初から \(K\) の倍数であるお菓子（\(A_i \bmod K = 0\)）は、選んでも選ばなくても「総和を \(K\) で割った余り」に一切影響を与えません。 このようなお菓子の個数を \(C_0\) とすると、これら \(C_0\) 個のお菓子の選び方は \(2^{C_0}\) 通りあります。

したがって、余りが \(0\) でないお菓子（\(A_i \bmod K \neq 0\)）だけでDPを行い、最終的に得られた「余りが \(0\) になる選び方の数」に \(2^{C_0}\) を掛ければ、全体での選び方を効率よく計算できます。これにより、DPのループ回数を \(N\) から \(N - C_0\) に減らすことができます。

アルゴリズム

1. コーナーケースの処理: \(K = 1\) の場合は、\((2^N - 1) \bmod 998244353\) を出力して終了します。
2. お菓子の分類: 各 \(A_i\) について、 \(A_i \bmod K\) を計算します。
   • 余りが \(0\) のものの個数を \(C_0\) とする。
   • 余りが \(0\) でないものをリスト A_filtered に追加する。
3. DPテーブルの定義と初期化: 長さ \(K\) の配列 dp を用意します。
   • dp[j] : A_filtered のお菓子からいくつか選んで、余りの総和が \(j\) になる選び方の数。
   • 初期状態：何も選んでいない状態（総和 \(0\)）が \(1\) 通りあるので、dp[0] = 1、それ以外は 0 とします。
4. DPの遷移: A_filtered の各要素 \(a\) について、DPテーブルを更新します。 $\(ndp[i] = (dp[i] + dp[(i - a) \bmod K]) \bmod 998244353\)$ これを省メモリ化のために、2つの配列を使い回しながら更新します。
5. 答えの計算: DP終了後、dp[0] は「余りが \(0\) でないお菓子からいくつか選んで、総和が \(K\) の倍数になる選び方の数」を表します（これには「1個もお菓子を選ばない」という \(1\) 通りも含みます）。 これに、余りが \(0\) のお菓子の選び方 \(2^{C_0}\) を掛け合わせます。 最後に、お菓子を1個も選ばない（空集合の）場合を引くために \(1\) を引きます。 $\(\text{Answer} = (dp[0] \times 2^{C_0} - 1) \bmod 998244353\)$

計算量

• 時間計算量: \(O(N + K + (N - C_0) \times K)\) お菓子の分類に \(O(N)\)、DPの遷移に最大で \(O(NK)\) かかります。制約より \(N \times K \le 2 \times 10^7\) であるため、実行時間制限に十分間に合います。
• 空間計算量: \(O(N + K)\) 入力の保存に \(O(N)\)、DPテーブルの保持に \(O(K)\) のメモリを使用します。

実装のポイント

• Pythonの負のインデックスの活用: 遷移式における dp[i - a] について、i - a が負の値になることがあります。Pythonではリストの負のインデックス list[-x] は末尾から x 番目の要素を指すため、dp[i - a] と書くだけで自動的に \(\bmod K\) の位置（正確には i - a + K）を参照してくれます。これにより、余りの計算を簡潔に書くことができます。

• 空集合の除外: 問題文では「空でない部分集合」を求めるよう指示されています。DPの初期値 dp[0] = 1 は「お菓子を1個も選ばない（総和0）」という空集合を含んでいるため、最後に必ず \(1\) を引く必要があります。また、引いた結果が負にならないよう、MODを加算してから余りを取るようにします。

  ソースコード

import sys

def solve():
    input = sys.stdin.read
    data = input().split()
    if not data:
        return
    N = int(data[0])
    K = int(data[1])
    A = [int(x) for x in data[2:]]
    
    MOD = 998244353
    
    if K == 1:
        ans = (pow(2, N, MOD) - 1 + MOD) % MOD
        print(ans)
        return

    C_0 = 0
    A_filtered = []
    for a in A:
        rem = a % K
        if rem == 0:
            C_0 += 1
        else:
            A_filtered.append(rem)
            
    dp = [0] * K
    dp[0] = 1
    ndp = [0] * K
    
    for a in A_filtered:
        for i in range(K):
            val = dp[i] + dp[i - a]
            if val >= MOD:
                val -= MOD
            ndp[i] = val
        dp, ndp = ndp, dp
        
    ans = (dp[0] * pow(2, C_0, MOD) - 1 + MOD) % MOD
    print(ans)

if __name__ == '__main__':
    solve()

---

この解説は gemini-3.5-flash-high によって生成されました。