C - お菓子の詰め合わせ / Assortment of Sweets 解説 by admin
gpt-5.5-highOverview
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
- For each \(A_i\), calculate \(A_i \bmod K\) and count the number of occurrences for each remainder.
- Set \(dp[0] = 1\) and all other elements to \(0\).
- 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] \)$
- 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 \)$
- 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.