2228E1UnreviewedunratedGPT-5.5 (high)

Amanojaku and Sequence (Easy Version)

Progressive hints first, then the full explanation and implementation when you're ready to cash out.

combinatorics math

Original problem

Review status

AI-generated and still unreviewed. Double-check the details before internalizing them.

Hints

Progressive nudges

Open only as much as you need to keep the solve alive.

Forget updates and multiple queries — easy version is just one static subarray. First ask: after fixing all non- $-1$ values, how much total mass is left for the wildcards?

Let the subarray contain $k$ wildcards and let their remaining sum be $M$ . If $M<0$ , you're dead: no valid sequence. Otherwise every filling is just a composition of $M$ into $k$ non-negative parts.

For a position $i$ , split its prefix sum into two parts:

\text{prefix}_i = F_i + X_i,

where $F_i$ is the sum of fixed values up to $i$ , and $X_i$ is the sum of wildcard variables up to $i$ .

If $t_i$ wildcards appear up to position $i$ , then $X_i$ is the sum of the first $t_i$ variables in a uniformly chosen composition of $M$ into $k$ parts. You only need three aggregated sums over positions:

\sum F_i^2,\qquad \sum F_i t_i,\qquad \sum t_i,\qquad \sum t_i^2.

The number of compositions is

N=\binom{M+k-1}{k-1}.

For $X_t=$ sum of first $t$ parts:

\sum X_t = N\cdot \frac{Mt}{k},

and

\sum X_t^2 = N\left(\frac{M t(k-t)(M+k)}{k^2(k+1)}+\frac{M^2t^2}{k^2}\right).

Plug this into $(F_i+X_{t_i})^2$ and sum over all positions.

Main reduction

We need compute $g(b,m)$ for one subarray $b=[a_l,\dots,a_r]$ .

Let:

$k$ = number of $-1$ positions in the subarray,
$S$ = sum of fixed non-negative values in the subarray,
$M=m-S$ = total sum that must be distributed among wildcard positions.

If $M<0$ , there are no valid sequences, so the answer is $0$ .

If $k=0$ , there is only one possible sequence. It is valid iff $M=0$ , and the answer is just

\sum_i \left(\sum_{j\le i} b_j\right)^2.

Now assume $k>0$ .

Each valid sequence corresponds exactly to a composition

x_1+x_2+\cdots+x_k=M,\qquad x_i\ge 0.

The number of such compositions is

N=\binom{M+k-1}{k-1}.

Classic stars and bars. No magic, no goblins.

Prefix structure

For every position $i$ in the subarray, define:

$F_i$ = sum of fixed values up to position $i$ ,
$t_i$ = number of wildcards up to position $i$ ,
$X_{t_i}=x_1+x_2+\cdots+x_{t_i}$ .

Then the actual prefix sum at position $i$ is

F_i+X_{t_i}.

So the contribution of position $i$ over all valid fillings is

\sum_{\text{compositions}} (F_i+X_{t_i})^2.

Expand the square:

\sum (F_i+X_{t_i})^2 = N F_i^2 + 2F_i\sum X_{t_i}+\sum X_{t_i}^2.

So we only need the first two moments of $X_t$ , where $X_t$ is the sum of the first $t$ parts in a uniformly chosen composition of $M$ into $k$ parts.

Moment formulas

For $X_t=x$ , the number of compositions with that value is

\binom{x+t-1}{t-1}\binom{M-x+k-t-1}{k-t-1},

with the natural edge cases. This is the beta-binomial distribution with parameters $M,t,k-t$ .

The useful formulas are:

\sum X_t = N\cdot \frac{Mt}{k},

and

\sum X_t^2 = N\left( \frac{M t(k-t)(M+k)}{k^2(k+1)} +\frac{M^2t^2}{k^2} \right).

The first one is also obvious by symmetry: total mass $M$ is shared equally among $k$ variables on average.

The second one is variance plus mean squared:

\operatorname{Var}(X_t) = \frac{M t(k-t)(M+k)}{k^2(k+1)}, \qquad \mathbb{E}[X_t]=\frac{Mt}{k}.

All division is modulo $998244353$ . Since $k\le 3\cdot 10^5 < 998244353$ , denominators are invertible. Nice.

Aggregating over positions

For every position, we need $F_i$ and $t_i$ . While scanning the subarray once, compute:

A_0=\sum_i F_i^2,

A_1=\sum_i F_i t_i,

T_1=\sum_i t_i,

T_2=\sum_i t_i^2.

Now sum the expanded formula over all positions.

The answer is:

N A_0 +2\left(N\frac{M}{k}\right)A_1 +N\left( \frac{M(M+k)}{k^2(k+1)}\sum_i t_i(k-t_i) +\frac{M^2}{k^2}\sum_i t_i^2 \right).

And

\sum_i t_i(k-t_i)=kT_1-T_2.

So everything is determined by those four simple aggregates. Beautifully unfair for anyone trying DP.

Complexity

For each test case, we scan the queried subarray a constant number of times. Since $q=1$ and the sum of $n$ over test cases is bounded,

O\left(\sum n + M_{\max}+N_{\max}\right)

time is enough, mainly for factorial precomputation and array scanning.

Memory usage is

O(10^6+3\cdot 10^5)

for factorials and inverse factorials.

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

using ll = long long;

void setIO() {
    ios::sync_with_stdio(false);
    cin.tie(nullptr);
}

const ll MOD = 998244353;
const int MAXF = 1300005; // max m + max n + a little safety

ll modpow(ll a, ll e) {
    ll r = 1;
    while (e) {
        if (e & 1) r = r * a % MOD;
        a = a * a % MOD;
        e >>= 1;
    }
    return r;
}

vector<ll> fact, ifact;

ll C(int n, int r) {
    if (r < 0 || r > n || n < 0) return 0;
    return fact[n] * ifact[r] % MOD * ifact[n - r] % MOD;
}

int main() {
    setIO();

    fact.assign(MAXF, 1);
    ifact.assign(MAXF, 1);
    for (int i = 1; i < MAXF; i++) fact[i] = fact[i - 1] * i % MOD;
    ifact[MAXF - 1] = modpow(fact[MAXF - 1], MOD - 2);
    for (int i = MAXF - 2; i >= 0; i--) ifact[i] = ifact[i + 1] * (i + 1) % MOD;

    int TC;
    cin >> TC;
    while (TC--) {
        int n, q;
        cin >> n >> q;

        vector<ll> a(n + 1);
        for (int i = 1; i <= n; i++) cin >> a[i];

        while (q--) {
            int op, l, r;
            ll m;
            cin >> op >> l >> r >> m;

            ll fixedSum = 0;
            int k = 0;
            for (int i = l; i <= r; i++) {
                if (a[i] == -1) k++;
                else fixedSum += a[i];
            }

            if (fixedSum > m) {
                cout << 0 << '\n';
                continue;
            }

            ll M = m - fixedSum;

            if (k == 0) {
                if (M != 0) {
                    cout << 0 << '\n';
                    continue;
                }

                ll pref = 0, ans = 0;
                for (int i = l; i <= r; i++) {
                    pref = (pref + a[i]) % MOD;
                    ans = (ans + pref * pref) % MOD;
                }
                cout << ans << '\n';
                continue;
            }

            ll N = C((int)(M + k - 1), k - 1);

            ll prefFixed = 0;
            int wildSeen = 0;

            ll A0 = 0; // sum F_i^2
            ll A1 = 0; // sum F_i * t_i
            ll T1 = 0; // sum t_i
            ll T2 = 0; // sum t_i^2

            for (int i = l; i <= r; i++) {
                if (a[i] == -1) {
                    wildSeen++;
                } else {
                    prefFixed += a[i];
                }

                ll F = prefFixed % MOD;
                ll T = wildSeen % MOD;

                A0 = (A0 + F * F) % MOD;
                A1 = (A1 + F * T) % MOD;
                T1 = (T1 + T) % MOD;
                T2 = (T2 + T * T) % MOD;
            }

            ll kMod = k % MOD;
            ll MMod = M % MOD;
            ll invK = modpow(kMod, MOD - 2);
            ll invK2 = invK * invK % MOD;
            ll invK1 = modpow((k + 1) % MOD, MOD - 2);

            ll ans = 0;

            // N * sum F_i^2
            ans = (ans + N * A0) % MOD;

            // 2 * (N * M / k) * sum F_i t_i
            ll linearCoeff = N * MMod % MOD * invK % MOD;
            ans = (ans + 2 * linearCoeff % MOD * A1) % MOD;

            // N * [ M(M+k)/(k^2(k+1)) * sum t_i(k-t_i)
            //       + M^2/k^2 * sum t_i^2 ]
            ll sumTkMinusT2 = (kMod * T1 % MOD - T2 + MOD) % MOD;

            ll coefVar = MMod * ((M + k) % MOD) % MOD;
            coefVar = coefVar * invK2 % MOD * invK1 % MOD;

            ll coefMeanSq = MMod * MMod % MOD * invK2 % MOD;

            ll secondMomentPart = (coefVar * sumTkMinusT2 + coefMeanSq * T2) % MOD;
            ans = (ans + N * secondMomentPart) % MOD;

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

    return 0;
}

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

using ll = long long;

void setIO() {
    ios::sync_with_stdio(false);
    cin.tie(nullptr);
}

const ll MOD = 998244353;
const int MAXF = 1300005; // max m + max n + a little safety

ll modpow(ll a, ll e) {
    ll r = 1;
    while (e) {
        if (e & 1) r = r * a % MOD;
        a = a * a % MOD;
        e >>= 1;
    }
    return r;
}

vector<ll> fact, ifact;

ll C(int n, int r) {
    if (r < 0 || r > n || n < 0) return 0;
    return fact[n] * ifact[r] % MOD * ifact[n - r] % MOD;
}

int main() {
    setIO();

    fact.assign(MAXF, 1);
    ifact.assign(MAXF, 1);
    for (int i = 1; i < MAXF; i++) fact[i] = fact[i - 1] * i % MOD;
    ifact[MAXF - 1] = modpow(fact[MAXF - 1], MOD - 2);
    for (int i = MAXF - 2; i >= 0; i--) ifact[i] = ifact[i + 1] * (i + 1) % MOD;

    int TC;
    cin >> TC;
    while (TC--) {
        int n, q;
        cin >> n >> q;

        vector<ll> a(n + 1);
        for (int i = 1; i <= n; i++) cin >> a[i];

        while (q--) {
            int op, l, r;
            ll m;
            cin >> op >> l >> r >> m;

            ll fixedSum = 0;
            int k = 0;
            for (int i = l; i <= r; i++) {
                if (a[i] == -1) k++;
                else fixedSum += a[i];
            }

            if (fixedSum > m) {
                cout << 0 << '\n';
                continue;
            }

            ll M = m - fixedSum;

            if (k == 0) {
                if (M != 0) {
                    cout << 0 << '\n';
                    continue;
                }

                ll pref = 0, ans = 0;
                for (int i = l; i <= r; i++) {
                    pref = (pref + a[i]) % MOD;
                    ans = (ans + pref * pref) % MOD;
                }
                cout << ans << '\n';
                continue;
            }

            ll N = C((int)(M + k - 1), k - 1);

            ll prefFixed = 0;
            int wildSeen = 0;

            ll A0 = 0; // sum F_i^2
            ll A1 = 0; // sum F_i * t_i
            ll T1 = 0; // sum t_i
            ll T2 = 0; // sum t_i^2

            for (int i = l; i <= r; i++) {
                if (a[i] == -1) {
                    wildSeen++;
                } else {
                    prefFixed += a[i];
                }

                ll F = prefFixed % MOD;
                ll T = wildSeen % MOD;

                A0 = (A0 + F * F) % MOD;
                A1 = (A1 + F * T) % MOD;
                T1 = (T1 + T) % MOD;
                T2 = (T2 + T * T) % MOD;
            }

            ll kMod = k % MOD;
            ll MMod = M % MOD;
            ll invK = modpow(kMod, MOD - 2);
            ll invK2 = invK * invK % MOD;
            ll invK1 = modpow((k + 1) % MOD, MOD - 2);

            ll ans = 0;

            // N * sum F_i^2
            ans = (ans + N * A0) % MOD;

            // 2 * (N * M / k) * sum F_i t_i
            ll linearCoeff = N * MMod % MOD * invK % MOD;
            ans = (ans + 2 * linearCoeff % MOD * A1) % MOD;

            // N * [ M(M+k)/(k^2(k+1)) * sum t_i(k-t_i)
            //       + M^2/k^2 * sum t_i^2 ]
            ll sumTkMinusT2 = (kMod * T1 % MOD - T2 + MOD) % MOD;

            ll coefVar = MMod * ((M + k) % MOD) % MOD;
            coefVar = coefVar * invK2 % MOD * invK1 % MOD;

            ll coefMeanSq = MMod * MMod % MOD * invK2 % MOD;

            ll secondMomentPart = (coefVar * sumTkMinusT2 + coefMeanSq * T2) % MOD;
            ans = (ans + N * secondMomentPart) % MOD;

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

    return 0;
}

View on Codeforces Back to the list