Summary

This problem asks us to find the number of ways to choose one 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 try to track the sum of calories directly, the values will become extremely large. However, the condition “the sum is a multiple of \(K\)” can be rephrased as “the sum modulo \(K\) is \(0\)”. Therefore, we only need to consider the calorie of each candy modulo \(K\).

If we try to brute-force all possible subsets of candies (bit brute-force), there are \(2^N\) ways to choose, which will not run within the time limit under the constraint \(N \le 2 \times 10^5\).

Thus, we consider Dynamic Programming (DP). We can calculate this efficiently by keeping track of “the number of subsets of the candies processed so far whose sum modulo \(K\) is \(r\)”. The number of states is \(K\), and for each candy, we consider \(K\) transitions, resulting in an overall time complexity of \(O(NK)\). Given the constraint \(N \times K \le 2 \times 10^7\), this complexity is well within the time limit.

Algorithm

We solve this using Dynamic Programming (DP).

1. Convert the calorie of each candy \(A_i\) to its remainder modulo \(K\).
2. Prepare a 1D array dp of length \(K\). dp[r] represents the number of ways to choose candies so far such that their sum modulo \(K\) is \(r\).
3. As the initial state, since choosing nothing results in a sum of \(0\), we set dp[0] = 1.
4. Process the candies one by one. Let \(a\) be the calorie of the current candy modulo \(K\).
   • If \(a = 0\): Choosing or not choosing this candy does not change the remainder. Therefore, for every element in dp, there are 2 choices (“choose” and “do not choose”), so we double each element.
   • If \(a \neq 0\): Prepare a new array nxt and calculate the transitions for either choosing or not choosing the current candy. When transitioning from the remainder j, choosing the candy changes the remainder to (j + a) % K. Therefore, to the (j + a) % K-th element of nxt, we add the sum of “the original value of not choosing, dp[(j + a) % K]” and “the value of choosing, dp[j]”.
5. After processing all candies, dp[0] will be the number of ways to choose candies such that their sum is divisible by \(K\). However, this includes the 1 way of choosing nothing (the empty set) in addition to those that satisfy the “choose one or more” condition. Therefore, the final answer is (dp[0] - 1) % MOD.

Complexity

• Time Complexity: \(O(NK)\)
• Space Complexity: \(O(K)\)

Implementation Details

• If there are consecutive candies with a remainder of \(0\), running a loop of length \(K\) every time would be wasteful. Thus, we speed up the computation by handling \(a = 0\) as a special case where we simply double each element of the array.

• In the transition for a != 0, instead of writing nxt[j] = dp[j] + dp[(j - a + K) % K], the code updates it as nxt[(j + a) % K] = dp[(j + a) % K] + dp[j]. This allows us to add both the “not choosing” and “choosing” transitions to all elements at once by simply scanning j from \(0\) to \(K-1\).

• To prevent the result of the subtraction dp[0] - 1 (to exclude the empty set) from becoming negative, remember to compute (dp[0] - 1) % MOD properly (e.g., (dp[0] - 1 + MOD) % MOD in languages where % can return negative values).

  Source Code

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()

This editorial was generated by or-glm5.2-high.