Codeforces
CF Step
Youtube Linkedin Discord Toggle Dark/Light/Auto mode Toggle Dark/Light/Auto mode Toggle Dark/Light/Auto mode Back to homepage

Code : Colored Balls

Expression DP : Quadratic

#include <atcoder/modint>
#include <bits/stdc++.h>

using namespace std;
using namespace atcoder;

using mint = modint998244353;

int halfceil(int num) { return (num / 2 + (num % 2 != 0)); }

void solve(vector<int> &a) {
    int n = a.size();
    int mx = accumulate(a.begin(), a.end(), 0) + 2;

    vector<mint> dp(mx, 0);
    // dp[sum] is the number of subsets with sum = sum and max = a[i].
    // All such subsets should necessarily include i-th element.

    sort(a.begin(), a.end());

    dp[0] = 1;
    mint ans = 0;
    for (int i = 0; i < n; i++) {
        // ndp means the subsets necessarily ending at element a[i].
        vector<mint> ndp(mx, 0);
        for (int sum = 0; sum < mx; sum++) {
            if (dp[sum] == 0) {
                continue;
            }
            ndp[sum + a[i]] = dp[sum];
        }
        // We can't store the max value in our DP, so we have to accumulate
        // the answer immediately, as we know that all DP values necessarily
        // contain the current value as max.
        for (int sum = 0; sum < mx; sum++) {
            ans += max(a[i], halfceil(sum)) * ndp[sum];
        }

        // We combine subsets which aren't necessarily ending at a[i].
        // This will needed for future subset construction.
        for (int sum = 0; sum < mx; sum++) {
            dp[sum] += ndp[sum];
        }
    }

    cout << ans.val() << "\n";
}

int main() {
    ios_base::sync_with_stdio(false);
    cin.tie(NULL);

    int n;
    cin >> n;
    vector<int> a(n);
    for (int i = 0; i < n; i++) {
        cin >> a[i];
    }
    solve(a);
    return 0;
}

Expression DP : Cubic

#include <atcoder/modint>
#include <bits/stdc++.h>

using namespace std;
using namespace atcoder;

using mint = modint998244353;

int halfceil(int num) { return (num / 2 + (num % 2 != 0)); }

void solve(vector<int> &a) {
    int n = a.size();
    int mx = accumulate(a.begin(), a.end(), 0) + 2;

    vector<vector<mint>> dp(mx, vector<mint>(mx, 0));
    // dp[m][sum] is the number of subsets with max = m and sum = sum.

    dp[0][0] = 1;
    for (int i = 0; i < n; i++) {
        auto ndp = dp;
        for (int m = 0; m < mx; m++) {
            for (int sum = 0; sum < mx; sum++) {
                if (dp[m][sum] == 0) {
                    continue;
                }
                ndp[max(m, a[i])][sum + a[i]] += dp[m][sum];
            }
        }
        swap(dp, ndp);
    }

    mint ans = 0;
    for (int m = 0; m < mx; m++) {
        for (int sum = 0; sum < mx; sum++) {
            ans += max(m, halfceil(sum)) * dp[m][sum];
        }
    }
    cout << ans.val() << "\n";
}

int main() {
    ios_base::sync_with_stdio(false);
    cin.tie(NULL);

    int n;
    cin >> n;
    vector<int> a(n);
    for (int i = 0; i < n; i++) {
        cin >> a[i];
    }
    solve(a);
    return 0;
}

Slots DP : Cubic

#include <atcoder/modint>
#include <bits/stdc++.h>

using namespace std;
using namespace atcoder;

using mint = modint998244353;

void solve(vector<int> &a) {
    int n = a.size();
    int mx = accumulate(a.begin(), a.end(), 0) + 2;

    vector<vector<mint>> dp(mx, vector<mint>(mx, 0));
    // dp[one][two] is the number of subsets where the count of isolated slots
    // is equal to "one" and the count of paired slots is equal to "two".

    dp[0][0] = 1;
    for (int i = 0; i < n; i++) {
        auto ndp = dp;
        for (int one = 0; one < mx; one++) {
            for (int two = 0; two < mx; two++) {
                if (dp[one][two] == 0) {
                    continue;
                }

                int none = one, ntwo = two, backup = a[i];
                // First, pair a[i] with isolated elements.
                int slots = min(a[i], one);
                a[i] -= slots;
                none -= slots;
                ntwo += slots;

                // Next, correct the pair (x, y) and (a[i], a[i]) to
                // (x, a[i]) and (y, a[i]).
                slots = min(a[i] / 2, two);
                a[i] -= 2 * slots;
                ntwo += slots;

                // The remaining a[i] will form its own isolated slots.
                none += a[i];

                a[i] = backup;
                ndp[none][ntwo] += dp[one][two];
            }
        }
        swap(dp, ndp);
    }

    mint ans = 0;
    for (int one = 0; one < mx; one++) {
        for (int two = 0; two < mx; two++) {
            ans += (one + two) * dp[one][two];
        }
    }
    cout << ans.val() << "\n";
}

int main() {
    ios_base::sync_with_stdio(false);
    cin.tie(NULL);

    int n;
    cin >> n;
    vector<int> a(n);
    for (int i = 0; i < n; i++) {
        cin >> a[i];
    }
    solve(a);
    return 0;
}