公式

C - お菓子の詰め合わせ / Assortment of Sweets 解説 by MMNMM


\(1\) つもお菓子を選ばないような選び方もあわせて考えます(\(1\) つもお菓子を選ばないときの合計カロリーは \(0\) で、これは \(K\) で割り切れます。よって、求めるべき答えに \(1\) を足したものが求められます)。

\(\operatorname{dp} _ i[j]\coloneqq i\) 番目までのお菓子を選ぶ方法のうち、選んだお菓子のカロリーの総和を \(K\) で割った余りが \(j\) であるような選び方の総数 \((0\le i\le N,0\le j\lt K)\) と定めます。 すると、\(i\) 番目のお菓子を選ぶか否かで場合分けをすることで、次のような漸化式を得ることができます。

\[\operatorname{dp} _ 0[j]=\begin{cases}1&(j=0)\\0&(0\lt j\lt K)\end{cases}\] \[\operatorname{dp} _ i[j]=\operatorname{dp} _ {i-1}[j]+\operatorname{dp} _ {i-1}[(j-A _ i)\,\%\,K]\ (1\le i\le N)\]

ここで、\(x\,\%\,y\) は \(x\) を \(y\) で割った余りを表すとします。

あとは、これに従って DP テーブルを計算することで時間計算量を \(O(NK)\) とすることができます。 \(\operatorname{dp} _ i\) の計算には \(\operatorname{dp} _ {i-1}\) しか必要ないため、適切に実装することで空間計算量を \(O(K)\) にすることもできます。

実装例は以下のようになります。

#include <iostream>
#include <vector>
#include <atcoder/modint>
using namespace std;
using modint = atcoder::static_modint<998244353>;

int main() {
    int N, K;
    cin >> N >> K;
    vector<modint> dp(K), prev(K);
    dp[0] = 1;
    for (int i = 0; i < N; ++i) {
        int A;
        cin >> A;
        swap(dp, prev);
        for (int j = 0; j < K; ++j) {
            dp[j] = prev[j] + prev[(j + K - A % K) % K];
        }
    }
    cout << dp[0].val() - 1 << endl;
    return 0;
}
N, K = map(int, input().split())

dp = [0 for i in range(K)]
prev = [0 for i in range(K)]
dp[0] = 1

for A in map(int, input().split()):
    dp, prev = prev, dp
    for j in range(K):
        dp[j] = (prev[j] + prev[(j - A) % K]) % 998244353

print(dp[0] - 1)