CF1954D Colored Balls 题解

原题链接

思路简介

对于一组可能的小球组合，设数量和为 $sum$，最大值为 $a$，那么能分成的最小组数即为 $\min (a,\lceil \frac{sum}{2}\rceil )$。观察到 $sum$ 的值不会大于 $5000$，我们可以用一个 $cnt$ 数组记录每一种数量和的出现次数，将数量从小到大排序后一边通过遍历 $cnt$ 数组分情况来统计答案，一边更新 $cnt$ 数组即可。

代码

#include <bits/stdc++.h>
#define read std::cin
#define write std::cout
typedef long long ll;

const int N = 5005;
const int mod = 998244353;

int n, a[N];
ll dp[N], res;

int main() {
    std::ios::sync_with_stdio(false);
    read.tie(0), write.tie(0);
    read >> n;
    for (int i = 1; i <= n; ++i) read >> a[i];
    std::sort(a + 1, a + 1 + n);
    dp[0] = 1ll; // 0 的出现次数为 1
    for (int i = 1; i <= n; ++i) {
        for (int j = 5000; j >= a[i] + 1; --j) {
            // 每一组对答案的贡献为 (j + a[i] + 1) / 2
            res = (res + (j + a[i] + 1) / 2 * dp[j] % mod) % mod;
            if (j + a[i] <= 5000) dp[j + a[i]] = (dp[j + a[i]] + dp[j]) % mod; // 更新 cnt 数组
        }
        for (int j = a[i]; j >= 0; --j) {
            // 每一组对答案的贡献为 a[i]
            res = (res + dp[j] * a[i] % mod) % mod;
            if (j + a[i] <= 5000) dp[j + a[i]] = (dp[j + a[i]] + dp[j]) % mod;
        }
    }
    write << res;
    return 0;
}