Summary

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

Analysis

If we think about it naively, there are \(2\) choices (“choose” or “do not choose”) for each of the \(N\) candies, resulting in \(2^N\) possible subsets in total. Since \(N\) can be up to \(2 \times 10^5\), a method that exhaustively enumerates all subsets and checks if their sum is a multiple of \(K\) will not run in time (exponential time complexity).

The key observation here is that to determine whether the sum is divisible by \(K\), we do not need the exact sum itself; we only need to know the sum modulo \(K\).

For example, when \(K = 3\), a candy with \(7\) calories and a candy with \(4\) calories both have a remainder of \(1\) when divided by \(K\). In terms of whether the total sum becomes a multiple of \(K\), replacing each candy’s calories with its “remainder modulo \(K\)” does not change the answer.

Therefore, we can maintain the number of ways to choose candies using the sum of the chosen candies modulo \(K\) as the key. Since there are only \(K\) possible remainders (\(0, 1, \ldots, K-1\)), we can define a dynamic programming (DP) state based on this.

Algorithm

Let dp[j] be the “number of ways to choose a subset from the candies processed so far such that the sum modulo \(K\) is \(j\)”.

• Initial State: Since the remainder is \(0\) when we choose nothing (the empty set), we set dp[0] = 1 and all other states to \(0\).

Next, we process the candies one by one. Let the remainder of the \(i\)-th candy’s calories be \(r = A_i \bmod K\). For each candy, we have two choices: “do not choose” or “choose”.

The new state ndp[j] will be the sum of:

• If we do not choose this candy: The previous remainder was \(j\), which remains \(j\) \(\to\) dp[j]
• If we choose this candy: The previous remainder was \((j - r) \bmod K\), and adding \(r\) to it makes the remainder \(j\) \(\to\) dp[(j - r) \bmod K]

Thus:

\[ \text{ndp}[j] = \text{dp}[j] + \text{dp}[(j - r) \bmod K] \]

We repeat this process for all candies.

Finally, dp[0] represents the “number of ways to choose candies such that the sum is a multiple of \(K\)”. However, this includes the empty set (choosing nothing). Since the problem requires us to “choose \(1\) or more candies”, the answer is dp[0] - 1, subtracting \(1\) to account for the empty set.

Example

Let’s consider the case where \(N = 3\), \(K = 3\), and \(A = (1, 2, 3)\). The remainders are \(1, 2, 0\), respectively.

• Initial state: dp = [1, 0, 0]
• After processing Candy 1 (\(r=1\)): dp = [1, 1, 0] (empty set / \(\{1\}\))
• After processing Candy 2 (\(r=2\)): dp = [1, 1, 1] … The states accumulate: remainder 0 corresponds to \(\{\}\) and \(\{1,2\}\), remainder 1 corresponds to \(\{1\}\) and \(\{2\}\), etc.
• After processing Candy 3 (\(r=0\)): dp = [2, 2, 2]

The number of ways to get a remainder of \(0\) is dp[0] = 2 (the empty set and \(\{1,2\}\)). Excluding the empty set, the answer is \(1\). Indeed, the only non-empty subset whose sum is a multiple of \(3\) is \(\{1,2\}\) (sum is \(3\)), which is correct.

Complexity

• Time Complexity: \(O(N \times K)\)
  • This is because we update \(K\) remainder states for each candy. Since \(N \times K \leq 2 \times 10^7\) from the constraints, this is sufficiently fast.
• Space Complexity: \(O(K)\)
  • We only need to maintain two DP arrays (dp and ndp).

Implementation Details

• Be careful with negative remainders: When calculating \((j - r) \bmod K\), \(j - r\) can be negative. Since the % operator in C++ returns a negative value for negative inputs, we need to adjust it, for example, with if(prev < 0) prev += K;.

• Perform calculations modulo MOD: Since the answer can be very large, we take the modulo \(998244353\) after each addition. In the code, we use if(ndp[j] >= MOD) ndp[j] -= MOD; after addition, which is faster than using the modulo operator.

• Reusing arrays: By swapping pointers with swap(dp, ndp), we can avoid the cost of allocating a new array in each iteration. Since ndp is overwritten every time, there is no risk of forgetting to initialize it.

• Excluding the empty set: Do not forget to compute dp[0] - 1 at the end. To prevent the subtraction from resulting in a negative value, we compute it as (dp[0] - 1 + MOD) % MOD.

  Source Code

#include <bits/stdc++.h>
using namespace std;

const long long MOD = 998244353;

int main(){
    int N, K;
    scanf("%d %d", &N, &K);
    
    vector<long long> dp(K, 0), ndp(K, 0);
    dp[0] = 1; // empty subset
    
    for(int i = 0; i < N; i++){
        long long a;
        scanf("%lld", &a);
        int r = (int)(a % K);
        // ndp[j] = dp[j] + dp[(j - r) mod K]
        for(int j = 0; j < K; j++){
            int prev = j - r;
            if(prev < 0) prev += K;
            ndp[j] = dp[j] + dp[prev];
            if(ndp[j] >= MOD) ndp[j] -= MOD;
        }
        swap(dp, ndp);
    }
    
    long long ans = (dp[0] - 1 + MOD) % MOD;
    printf("%lld\n", ans);
    
    return 0;
}

---

This editorial was generated by claude4.8opus-high.