概要

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

考察

カロリーの総和そのものを管理しようとすると、値が非常に大きくなってしまいます。しかし、「総和が \(K\) の倍数になる」という条件は、「総和を \(K\) で割った余りが \(0\) になる」と言い換えることができます。したがって、各お菓子のカロリーを \(K\) で割った余りだけを考えればよいことがわかります。

お菓子の選び方を全探索（ビット全探索）しようとすると、選び方の総数が \(2^N\) 通りあり、制約の \(N \le 2 \times 10^5\) では計算が終わりません。

そこで動的計画法（DP）を考えます。「これまでに見たお菓子の部分集合で、総和を \(K\) で割った余りが \(r\) になるものの数」を管理することで、効率的に計算できます。状態の数は \(K\) 通りであり、お菓子1個につき \(K\) 通りの遷移を考えるため、全体の計算量は \(O(NK)\) となります。制約の \(N \times K \le 2 \times 10^7\) から、この計算量で十分間に合うことがわかります。

アルゴリズム

動的計画法（DP）を用いて解きます。

1. 各お菓子のカロリー \(A_i\) を \(K\) で割った余りに変換しておきます。
2. 長さ \(K\) の1次元配列 dp を用意します。dp[r] は「現時点で選んだお菓子の総和を \(K\) で割った余りが \(r\) になる選び方の数」を表します。
3. 初期状態として、何も選んでいない状態は総和 \(0\) なので dp[0] = 1 とします。
4. お菓子を1つずつ見ていきます。お菓子のカロリーを \(K\) で割った余りを \(a\) とします。
   • \(a = 0\) の場合: そのお菓子を選んでも選ばなくても余りは変わりません。よって、dp のすべての要素について、「選ぶ」と「選ばない」の2通りが存在するため、各要素を2倍します。
   • \(a \neq 0\) の場合: 新しい配列 nxt を用意し、現在のお菓子を選ぶ場合と選ばない場合の遷移を計算します。j 番目の余りから遷移するとき、お菓子を選ぶと余りは (j + a) % K になります。したがって、nxt の (j + a) % K 番目の要素には、「選ばない場合の元の値 dp[(j + a) % K]」と「選ぶ場合の値 dp[j]」の和を足します。
5. すべてのお菓子を処理し終わった後、dp[0] が「総和が \(K\) で割り切れる選び方の数」になります。ただし、これは「1個以上選ぶ」という条件を満たすものに加えて、「何も選ばない（空集合）」の1通りを含んでしまっているため、最終的な答えは (dp[0] - 1) % MOD とします。

計算量

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

実装のポイント

• 余りが \(0\) のお菓子が連続する場合、毎回長さ \(K\) のループを回すのは無駄になるため、\(a = 0\) のときは配列の各要素を2倍するという特別処理を行うことで計算を高速化しています。

• a != 0 の場合の遷移では、nxt[j] = dp[j] + dp[(j - a + K) % K] とする代わりに、コード内では nxt[(j + a) % K] = dp[(j + a) % K] + dp[j] という形で更新しています。これにより、j を \(0\) から \(K-1\) まで走査するだけで、すべての要素に「選ばない場合」と「選ぶ場合」の両方の遷移を一度に足すことができます。

• 空集合を除外するための引き算 dp[0] - 1 の結果が負になるのを防ぐため、最後に (dp[0] - 1) % MOD とすることを忘れないようにします。

  ソースコード

import sys

def solve():
    input = sys.stdin.readline
    N, K = map(int, input().split())
    MOD = 998244353
    
    A = [int(x) % K for x in input().split()]
    
    dp = [0] * K
    dp[0] = 1
    
    for a in A:
        if a == 0:
            dp = [(x * 2) % MOD for x in dp]
        else:
            nxt = [0] * K
            lim = K - a
            for j in range(lim):
                nxt[j + a] = (dp[j + a] + dp[j]) % MOD
            for j in range(lim, K):
                nxt[j + a - K] = (dp[j + a - K] + dp[j]) % MOD
            dp = nxt
            
    print((dp[0] - 1) % MOD)

solve()

この解説は or-glm5.2-high によって生成されました。