公式

C - お菓子の詰め合わせ / Assortment of Sweets 解説 by admin

gpt-5.5-high

Overview

We consider whether to choose each sweet or not, and find the number of non-empty subsets of sweets such that the sum of their calories is a multiple of \(K\).

The key point is that we only need to keep track of the “remainder modulo \(K\)” rather than the sum itself.

Analysis

There are \(2^N\) subsets in total.
Since \(N\) is at most \(2 \times 10^5\), any method that enumerates all subsets will definitely time out.

Therefore, for each sweet’s calorie \(A_i\), we focus only on

\(A_i \bmod K\)

instead of its actual value.

This is because whether the sum of a subset is divisible by \(K\) depends solely on whether the remainder of the sum is \(0\).

For example, when \(K = 5\):

  • A calorie of \(12\) has a remainder of \(2\)
  • A calorie of \(7\) also has a remainder of \(2\)

Thus, both can be treated as “sweets that increase the remainder of the sum by \(2\)”.


Now, let us consider the following DP:

\(dp[j] =\) the number of ways to choose some sweets from those processed so far such that the remainder of the sum modulo \(K\) is \(j\)

Initially, since there is a way to choose nothing, we have

\(dp[0] = 1\)

When we additionally consider one sweet with a remainder of \(r\), for each subset, we have two choices:

  • Do not choose this sweet
  • Choose this sweet

If the final remainder is \(j\), then:

  • If we do not choose it: the previous remainder must have been \(j\)
  • If we choose it: the previous remainder must have been \((j-r) \bmod K\)

Therefore, the transition is:

\(next[j] = dp[j] + dp[(j-r) \bmod K]\)


Also, sweets with a remainder of \(0\) are special.

Even if we choose a sweet with a remainder of \(0\), the remainder of the sum does not change.
Therefore, if there are \(c\) sweets with a remainder of \(0\), we can independently decide whether to choose each of them, so we can simply multiply the overall answer by \(2^c\).

Finally, since the DP includes the case of “choosing nothing”, we subtract \(1\) for the empty set.

Algorithm

  1. For each \(A_i\), calculate \(A_i \bmod K\) and count the number of occurrences for each remainder.
  2. Set \(dp[0] = 1\) and all other elements to \(0\).
  3. For each sweet with a remainder of \(1\) or greater, update the DP.
    • When processing one sweet with remainder \(r\), we set:

$\( next[j] = dp[j] + dp[(j-r) \bmod K] \)$

  1. Let \(cnt[0]\) be the number of sweets with a remainder of \(0\). Since they can be chosen freely without changing the remainder, the answer is:

$\( dp[0] \times 2^{cnt[0]} - 1 \)$

  1. Output the answer modulo \(998244353\).

As a concrete example, let’s consider the case where \(K=5\) and we process a sweet with a remainder of \(2\).

In order for the remainder to become \(3\) after choosing this sweet, the previous remainder must have been:

\[ (3 - 2) \bmod 5 = 1 \]

Therefore, we have:

\[ next[3] = dp[3] + dp[1] \]

Complexity

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

Since the constraints guarantee that \(N \times K \leq 2 \times 10^7\), this DP can be executed fast enough.

Implementation Details

  • Maintain the DP array with a length of \(K\).
  • In Python, writing old[j - r] will reference from the end of the list if \(j-r\) is negative. This perfectly matches the behavior of \(\bmod K\).
  • Sweets with a remainder of \(0\) are not processed in the DP; instead, we multiply by pow(2, cnt[0], MOD) at the end.
  • Since we need to exclude the empty set, we subtract \(1\) at the end.
  • If \(K=1\), the sum of any subset is always divisible by \(1\), so the answer is:

$\( 2^N - 1 \)$

Source Code

import sys

MOD = 998244353

def main():
    data = list(map(int, sys.stdin.buffer.read().split()))
    N, K = data[0], data[1]

    if K == 1:
        print((pow(2, N, MOD) - 1) % MOD)
        return

    cnt = [0] * K
    residues = []
    for a in data[2:]:
        r = a % K
        if cnt[r] == 0:
            residues.append(r)
        cnt[r] += 1

    dp = [0] * K
    dp[0] = 1

    m = MOD
    rk = range(K)

    for r in residues:
        if r == 0:
            continue
        c = cnt[r]
        for _ in range(c):
            old = dp
            dp = [(v if (v := old[j] + old[j - r]) < m else v - m) for j in rk]

    ans = (dp[0] * pow(2, cnt[0], m) - 1) % m
    print(ans)

if __name__ == "__main__":
    main()

This editorial was generated by gpt-5.5-high.