Official

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

GLM 5.2 (High)

Overview

This problem asks us to find the total 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\).

Intuition

Enumerating all ways to choose candies would take \(2^N\) ways, which is too slow since \(N\) can be up to \(2 \times 10^5\). Therefore, we consider counting them efficiently using Dynamic Programming (DP).

Managing the sum of the subset itself would make the values too large, but what the problem asks is whether the sum is a multiple of \(K\). In other words, we only need to keep track of the sum of the calories of the chosen candies modulo \(K\).

We maintain dp[j] as “the number of ways to choose candies from the first \(i\) candies such that the sum modulo \(K\) is \(j\)”. In the initial state, no candies are chosen, so the sum is \(0\), meaning there is \(1\) way to have a remainder of \(0\) (dp[0] = 1).

When processing the \(i\)-th candy (with calorie \(A_i\)), we have two choices: “choose” or “do not choose” this candy. - Do not choose: The remainder does not change. The current number of ways for remainder \(j\), dp[j], is carried over directly to the new next_dp[j]. - Choose: \(A_i \bmod K\) is added to the remainder \(j\). The new remainder becomes \((j + A_i \bmod K) \bmod K\). Therefore, we add the current number of ways dp[j] to next_dp[(j + A_i \bmod K) \bmod K].

Finally, after processing all candies, the number of ways to make the sum a multiple of \(K\) (remainder \(0\)) is stored in dp[0]. However, since the problem requires choosing “at least \(1\)” candy, we must subtract the state where nothing is chosen (the empty set, which is 1 way). Thus, the answer is dp[0] - 1.

Algorithm

  1. Initialize the DP table dp of size \(K\) with dp[0] = 1 and all other elements as \(0\).
  2. For each candy \(i\), repeat the following process:
    • Initialize a new DP table next_dp as a copy of dp (transition for when the candy is not chosen).
    • For each remainder \(j\) (\(0 \leq j < K\)), add dp[j] to next_dp[(j + A_i \bmod K) \bmod K] (transition for when the candy is chosen).
    • Update dp with next_dp.
  3. Finally, print dp[0] - 1. Note that since the answer should be modulo \(998244353\), you should adjust the calculation (e.g., add 998244353 before taking the modulo) to prevent the result from becoming negative.

Complexity

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

Given the constraints, \(N \times K \leq 2 \times 10^7\), so the \(O(N \times K)\) time complexity is well within the execution time limit.

Implementation Points

  • Although the provided code is written in C++, it has a Python file extension. In competitive programming, the execution speed of C++ is often advantageous, and for this problem where \(N \times K\) approaches \(2 \times 10^7\), using C++ is highly appropriate.

  • When updating the DP table, a new array next_dp is created each time. By copying the array, we can concisely represent the “do not choose” transition.

  • When calculating the remainder, setting idx = j + x and checking if idx >= K (then subtracting \(K\) if true) reduces the overhead of the % (modulo) operator, achieving a constant-factor speedup.

  • Finally, we output dp[0] - 1. Since dp[0] is never \(0\) as a mathematical value, we simply perform the subtraction and add MOD if the result becomes negative.

    Source Code

#include <iostream>
#include <vector>
#include <algorithm>

using namespace std;

int main() {
    ios_base::sync_with_stdio(false);
    cin.tie(NULL);
    int N, K;
    if (!(cin >> N >> K)) return 0;
    vector<int> A(N);
    for (int i = 0; i < N; ++i) {
        cin >> A[i];
    }
    
    const int MOD = 998244353;
    vector<int> dp(K, 0);
    vector<int> next_dp(K, 0);
    dp[0] = 1;
    
    for (int i = 0; i < N; ++i) {
        int x = A[i] % K;
        copy(dp.begin(), dp.end(), next_dp.begin());
        for (int j = 0; j < K; ++j) {
            int idx = j + x;
            if (idx >= K) idx -= K;
            next_dp[idx] += dp[j];
            if (next_dp[idx] >= MOD) next_dp[idx] -= MOD;
        }
        dp.swap(next_dp);
    }
    
    int ans = dp[0] - 1;
    if (ans < 0) ans += MOD;
    cout << ans << "\n";
    
    return 0;
}

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

posted:
last update: