Competitive Programming using Grok3

index.cpp

#include <bits/stdc++.h>
using namespace std;
using ll = long long;

const int MOD = 998244353;
const int MAXN = 5e5 + 10;

ll fac[MAXN], inv[MAXN];

// Modular exponentiation
ll pow_mod(ll a, ll b, ll mod) {
    ll res = 1;
    a %= mod;
    while (b) {
        if (b & 1) res = res * a % mod;
        a = a * a % mod;
        b >>= 1;
    }
    return res;
}

// Precompute factorials and inverses
void precompute() {
    fac[0] = 1;
    for (int i = 1; i < MAXN; i++) {
        fac[i] = fac[i-1] * i % MOD;
    }
    inv[MAXN-1] = pow_mod(fac[MAXN-1], MOD-2, MOD);
    for (int i = MAXN-2; i >= 0; i--) {
        inv[i] = inv[i+1] * (i+1) % MOD;
    }
}

// Compute multinomial coefficient: n! / (x_1! * x_2! * ...)
ll multinomial(int n, const vector<int>& x) {
    ll res = fac[n];
    for (int xi : x) {
        res = res * inv[xi] % MOD;
    }
    return res;
}

void solve() {
    vector<int> c(26);
    ll n = 0;
    for (int i = 0; i < 26; i++) {
        cin >> c[i];
        n += c[i];
    }
    if (n == 0) {
        cout << 1 << '\\\\n';
        return;
    }

    int odd_len = (n + 1) / 2; // ceil(n/2)
    int even_len = n / 2;       // floor(n/2)

    // dp[i][s]: ways to assign first i letters with sum of x_j = s
    vector<vector<ll>> dp(27, vector<ll>(odd_len + 1));
    dp[0][0] = 1;
    for (int i = 0; i < 26; i++) {
        for (int s = 0; s <= odd_len; s++) {
            if (dp[i][s] == 0) continue;
            // Assign all c[i] to odd positions
            if (s + c[i] <= odd_len) {
                dp[i+1][s + c[i]] = (dp[i+1][s + c[i]] + dp[i][s]) % MOD;
            }
            // Assign all c[i] to even positions
            dp[i+1][s] = (dp[i+1][s] + dp[i][s]) % MOD;
        }
    }

    ll ans = 0;
    // For each possible sum of x_i
    for (int s = 0; s <= odd_len; s++) {
        if (dp[26][s] == 0) continue;
        vector<int> x(26), y(26);
        int sum_y = 0;
        // Reconstruct x and y
        int curr_s = s;
        for (int i = 25; i >= 0; i--) {
            if (curr_s >= c[i] && dp[i][curr_s - c[i]] > 0) {
                x[i] = c[i];
                curr_s -= c[i];
            } else {
                y[i] = c[i];
                sum_y += c[i];
            }
        }
        if (sum_y != even_len) continue;
        // Compute multinomial coefficients
        ll ways_odd = multinomial(odd_len, x);
        ll ways_even = multinomial(even_len, y);
        // Total ways for this assignment
        ans = (ans + dp[26][s] * ways_odd % MOD * ways_even) % MOD;
    }

    cout << ans << '\\\\n';
}

int main() {
    ios::sync_with_stdio(false);
    cin.tie(nullptr);
    precompute();
    int t;
    cin >> t;
    while (t--) {
        solve();
    }
    return 0;
}