Overview

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\). Output the answer modulo \(998244353\).

Observation

Issues with the Naive Approach

Enumerating all subsets would result in \(2^N\) possibilities, which is far too slow for \(N \leq 2 \times 10^5\). An exponential-time approach is not feasible.

Key Observation

The key point to notice here is that to determine whether the sum of calories is divisible by \(K\), we only need to keep track of the sum modulo \(K\). Although the sum itself can become extremely large, there are only \(K\) possible remainders from \(0\) to \(K-1\).

Therefore, we can consider a dynamic programming (DP) approach where we maintain the state: “the number of ways to choose a subset of candies such that the sum of their calories modulo \(K\) is \(r\).”

Specifically, we will decide whether to “choose” or “not choose” each candy one by one, and update how many ways each remainder state can be achieved.

Algorithm

We define the array dp[r] as “the number of ways to choose candies from those processed so far such that the sum modulo \(K\) is \(r\).”

Initial State

In the state where no candies have been chosen yet (the empty set), the sum is \(0\). Thus, we set dp[0] = 1 and all other elements to \(0\).

Transitions

When processing the \(i\)-th candy (let \(a\) be its calorie value modulo \(K\)), there are two choices for each subset:

• Do not choose this candy: The remainder does not change.
• Choose this candy: The remainder increases by \(a\) (taken modulo \(K\)).

Therefore, the new dp is updated from the old dp as follows:

\[ \text{new\_dp}[(r + a) \bmod K] \mathrel{+}= \text{dp}[r] \]

Combining this with the “do not choose” case (which keeps the state as is), the new number of ways to have a remainder of \(r\) is the sum of these two contributions:

\[ \text{new\_dp}[r] = \text{dp}[r] + \text{dp}[(r - a) \bmod K] \]

• First term: The case where we did not choose the \(i\)-th candy (originally had a remainder of \(r\)).
• Second term: The case where we chose the \(i\)-th candy (had a remainder of \((r-a)\) before choosing it).

Special Case: \(a = 0\)

When \(a = 0\) (the calorie value is a multiple of \(K\)), choosing or not choosing the candy does not change the remainder. Therefore, the number of ways for all states is simply doubled.

Final Answer

After processing all candies, dp[0] will store “the number of ways to choose candies such that the sum is divisible by \(K\).” However, this includes the empty set (the case where we choose nothing), so the final answer is dp[0] - 1.

Optimization in Code

The second term \(\text{dp}[(r-a) \bmod K]\) corresponds to a cyclic shift (rotation) of the array by \(s = K - a\). In the code, we can do:

s = K - a
shifted = dp[s:] + dp[:s]

This makes the shifted array shifted[r] represent dp[(r-a) mod K], allowing us to batch-update the array as dp[r] + shifted[r].

Complexity

• Time Complexity: \(O(N \times K)\)
• Space Complexity: \(O(K)\)

Since updating the DP state for each candy takes \(O(K)\) time and we do this for \(N\) candies, the total time complexity is \(O(N \times K)\). Given the constraint \(N \times K \leq 2 \times 10^7\), this complexity is well within the time limit.

Implementation Details

• Although the calorie values \(A_i\) can be as large as \(10^9\), we only need \(A_i \bmod K\) for the DP. Thus, taking the modulo at the beginning is sufficient.

• Take the modulo \(\bmod\ 998244353\) at each update step to prevent overflow or the values from growing too large (in Python, this also prevents performance degradation caused by handling arbitrarily large integers).

• Do not forget to subtract \(1\) at the end to exclude the empty set. Calculating (dp[0] - 1) % MOD ensures that the result does not become negative even if dp[0] is \(0\).

  Source Code

import sys

def main():
    data = sys.stdin.buffer.read().split()
    N = int(data[0]); K = int(data[1])
    A = data[2:2+N]
    MOD = 998244353
    dp = [0] * K
    dp[0] = 1
    for x in A:
        a = int(x) % K
        if a == 0:
            dp = [(v * 2) % MOD for v in dp]
        else:
            s = K - a
            shifted = dp[s:] + dp[:s]
            dp = [(p + q) % MOD for p, q in zip(dp, shifted)]
    print((dp[0] - 1) % MOD)

main()

---

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