Overview

This problem asks us to find the number of ways to choose some (one or more) sweets from \(N\) sweets such that the sum of their calories is a multiple of \(K\). By using Dynamic Programming (DP), we can solve this problem quickly by leveraging the constraint \(N \times K \le 2 \times 10^7\).

Analysis

1. Focus on the Remainder of Calories

To determine whether the sum of the chosen sweets is divisible by \(K\), we do not need to deal with the exact calorie value \(A_i\) of each sweet. We are only interested in the remainder of the calories divided by \(K\), which is \(A_i \bmod K\). Therefore, we can simplify the problem by converting the calories of all sweets to their remainders modulo \(K\) beforehand.

2. Naive DP (Dynamic Programming)

Focusing on the remainder, we can consider the following DP: - \(dp[i][j]\): The number of ways to choose some sweets from the first \(i\) sweets such that the sum of their calories modulo \(K\) is \(j\).

We consider transitions based on whether we choose the \(i\)-th sweet (with remainder \(a = A_i \bmod K\)) or not: - Do not choose: \(dp[i-1][j]\) ways - Choose: \(dp[i-1][(j - a) \bmod K]\) ways

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

If we implement this DP naively, the number of states is \(N \times K\) and the transition takes \(O(1)\), resulting in an overall time complexity of \(O(NK)\). Since the constraint is \(N \times K \le 2 \times 10^7\), this complexity is well within the execution time limit.

3. Optimization for Further Speedup

In languages like Python, reducing the number of loop iterations as much as possible is crucial to avoid Time Limit Exceeded (TLE). Therefore, we apply the following two optimizations:

① Corner Case: \(K = 1\)

When \(K = 1\), the sum of any combination of sweets is always a multiple of \(1\) (i.e., any integer). Therefore, the number of non-empty subsets is simply \(2^N - 1\), which can be computed in \(O(\log N)\) without DP.

② Handling Sweets with Remainder \(0\) Separately

Sweets whose calories are already a multiple of \(K\) (\(A_i \bmod K = 0\)) do not affect the “sum modulo \(K\)” at all, whether they are chosen or not. If there are \(C_0\) such sweets, there are \(2^{C_0}\) ways to choose among them.

Thus, we can run the DP only on the sweets with non-zero remainders (\(A_i \bmod K \neq 0\)), and then multiply the resulting “number of ways to get a remainder of \(0\)” by \(2^{C_0}\) to efficiently find the total number of ways. This reduces the number of DP iterations from \(N\) to \(N - C_0\).

Algorithm

1. Handling Corner Cases: If \(K = 1\), output \((2^N - 1) \bmod 998244353\) and terminate.
2. Classifying the Sweets: For each \(A_i\), calculate \(A_i \bmod K\).
   • Let \(C_0\) be the number of sweets with a remainder of \(0\).
   • Add sweets with non-zero remainders to a list A_filtered.
3. DP Table Definition and Initialization: Prepare an array dp of length \(K\).
   • dp[j]: The number of ways to choose some sweets from A_filtered such that the sum of their remainders modulo \(K\) is \(j\).
   • Initial state: There is \(1\) way to choose nothing (sum \(0\)), so set dp[0] = 1 and all other elements to 0.
4. DP Transitions: For each element \(a\) in A_filtered, update the DP table: $\(ndp[i] = (dp[i] + dp[(i - a) \bmod K]) \bmod 998244353\)$ To save memory, we can update this in-place or reuse two arrays.
5. Calculating the Answer: After completing the DP, dp[0] represents the number of ways to choose some sweets from A_filtered such that their sum is a multiple of \(K\) (this includes the \(1\) way where no sweets are chosen). Multiply this by \(2^{C_0}\), which is the number of ways to choose sweets with a remainder of \(0\). Finally, subtract \(1\) to exclude the case where no sweets are chosen at all (the empty set). $\(\text{Answer} = (dp[0] \times 2^{C_0} - 1) \bmod 998244353\)$

Complexity

• Time Complexity: \(O(N + K + (N - C_0) \times K)\) Classifying the sweets takes \(O(N)\) time, and the DP transitions take at most \(O(NK)\) time. Since the constraint is \(N \times K \le 2 \times 10^7\), this easily runs within the time limit.
• Space Complexity: \(O(N + K)\) We use \(O(N)\) memory to store the input and \(O(K)\) memory to maintain the DP table.

Implementation Points

• Utilizing Negative Indexing in Python: In the transition formula dp[i - a], i - a can sometimes be negative. In Python, negative indexing list[-x] refers to the x-th element from the end of the list. Therefore, simply writing dp[i - a] automatically accesses the correct position modulo \(K\) (specifically, i - a + K). This allows us to write the remainder calculation very concisely.

• Excluding the Empty Set: The problem statement asks us to find “non-empty subsets”. Since the initial DP state dp[0] = 1 includes the empty set (choosing 0 sweets, sum 0), we must subtract \(1\) at the end. To prevent the result of subtraction from becoming negative, we should add the MOD before taking the modulo.

  Source Code

import sys

def solve():
    input = sys.stdin.read
    data = input().split()
    if not data:
        return
    N = int(data[0])
    K = int(data[1])
    A = [int(x) for x in data[2:]]
    
    MOD = 998244353
    
    if K == 1:
        ans = (pow(2, N, MOD) - 1 + MOD) % MOD
        print(ans)
        return

    C_0 = 0
    A_filtered = []
    for a in A:
        rem = a % K
        if rem == 0:
            C_0 += 1
        else:
            A_filtered.append(rem)
            
    dp = [0] * K
    dp[0] = 1
    ndp = [0] * K
    
    for a in A_filtered:
        for i in range(K):
            val = dp[i] + dp[i - a]
            if val >= MOD:
                val -= MOD
            ndp[i] = val
        dp, ndp = ndp, dp
        
    ans = (dp[0] * pow(2, C_0, MOD) - 1 + MOD) % MOD
    print(ans)

if __name__ == '__main__':
    solve()

---

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