Summary

This problem asks us to find the number of ways to choose some (one or more) of \(N\) candies such that the sum of their calories is a multiple of \(K\), modulo \(998244353\).

Analysis

1. Naive Approach and Its Limitations

For each candy, we have two choices: either to “choose” or “not choose” it. Thus, there are \(2^N\) possible ways in total. Under the constraint \(N \le 2 \times 10^5\), \(2^N\) is an extremely large number. Exploring all combinations (exhaustive search) will result in a Time Limit Exceeded (TLE) error.

2. Dynamic Programming (DP) Focusing on “Remainders”

The condition “the sum is a multiple of \(K\)” can be rephrased as “the remainder of the sum divided by \(K\) is \(0\)”. When we decide whether to choose each candy one by one, we do not need to keep track of the “exact sum” of the candies chosen so far. As long as we maintain “the sum modulo \(K\)”, we can correctly calculate the new remainder when the next candy is added.

By grouping the states by their “remainders” in this way, we can apply Dynamic Programming (DP).

3. DP Design

We define the DP table as follows:

• \(dp[i][j]\) : The number of ways to choose some candies from the first \(i\) candies such that the sum of their calories modulo \(K\) is \(j\).

The remainder can be any of the \(K\) values from \(0\) to \(K-1\). When considering the \((i+1)\)-th candy (with calorie \(A_{i+1}\) and remainder \(rem = A_{i+1} \bmod K\)), the transitions are the sum of the following two patterns:

1. When we do not choose the \((i+1)\)-th candy The remainder does not change. This transitions from the state where the remainder was \(j\) up to the \(i\)-th candy.
2. When we choose the \((i+1)\)-th candy If we let the remainder before choosing this candy be \(prev\), we must have \(prev + rem \equiv j \pmod K\). Solving this for \(prev\) gives \(prev \equiv j - rem \pmod K\).

Therefore, the transition formula is as follows: $\(dp[i+1][j] = (dp[i][j] + dp[i][(j - rem + K) \bmod K]) \bmod 998244353\)$

For the initial state, there is exactly \(1\) way to choose no candies at all (sum is \(0\), which means the remainder is \(0\)). Thus, we initialize it as follows: * \(dp[0][0] = 1\) * \(dp[0][j] = 0 \quad (1 \le j < K)\)

4. Excluding the Empty Set

We want to find the number of ways to choose one or more candies. The DP calculation includes the \(1\) way of choosing no candies at all (the empty set, which has a sum of \(0\) and remainder \(0\)). Therefore, the final answer is \(dp[N][0] - 1\).

Algorithm

1. Prepare an array dp of length \(K\), and initialize dp[0] = 1 and all other elements to 0.
2. For each candy \(A_i\) (\(1 \le i \le N\)) in order, perform the following:
   • Find the remainder of \(A_i\) divided by \(K\): \(rem = A_i \bmod K\).
   • Prepare an array next_dp to represent the next state.
   • For each \(j\) (\(0 \le j < K\)), perform the following transition: next_dp[j] = (dp[j] + dp[(j - rem + K) % K]) % MOD
   • Update dp with next_dp.
3. Output (dp[0] - 1 + MOD) % MOD as the final answer.

Complexity

• Time Complexity: \(O(NK)\) For each of the \(N\) candies, we perform \(K\) transitions. Since \(N \times K \le 2 \times 10^7\) under the given constraints, this will easily run within the time limit.
• Space Complexity: \(O(K)\) Storing the 2D array \(dp[N+1][K]\) as is would consume a lot of memory. However, since the transitions only use the information from the immediately preceding row (the \(i\)-th state), we can reuse a 1D array to keep the space complexity down to \(O(K)\).

Implementation Details

• Memory Optimization: As shown in the C++ implementation code, we can significantly reduce memory usage by preparing two arrays of size \(K\) (dp and next_dp) and updating them alternately.

• Speed Optimization (Avoiding Modulo Arithmetic): During transitions, the % MOD operation is relatively slow. By replacing it with a subtraction when the sum exceeds MOD (e.g., if (val_next >= MOD) val_next -= MOD;), we can significantly reduce the execution time.

• Modulo of Negative Numbers: When calculating dp[0] - 1 at the end, if the value is \(0\), the result becomes \(-1\). To avoid negative remainders, make sure to add MOD before taking the modulo (i.e., (dp[0] - 1 + MOD) % MOD).

  Source Code

#include <iostream>
#include <vector>

using namespace std;

int main() {
    // 高速な入出力
    ios_base::sync_with_stdio(false);
    cin.tie(NULL);

    int N, K;
    if (!(cin >> N >> K)) return 0;

    const int MOD = 998244353;
    vector<int> dp(K, 0);
    dp[0] = 1; // 空集合（総和0）の1通り

    vector<int> next_dp(K);

    for (int i = 0; i < N; ++i) {
        long long val;
        cin >> val;
        int rem = val % K;
        for (int j = 0; j < K; ++j) {
            int prev = j - rem;
            if (prev < 0) prev += K;
            int val_next = dp[j] + dp[prev];
            if (val_next >= MOD) val_next -= MOD;
            next_dp[j] = val_next;
        }
        dp = next_dp;
    }

    // 空でない部分集合を求めるため、何も選ばない1通りを引く
    int ans = dp[0] - 1;
    if (ans < 0) ans += MOD;

    cout << ans << "\n";

    return 0;
}

---

This editorial was generated by gemini-3.5-flash-high.