Grenade throws from Valera#12198

Grenade throws from Valera (Editorial)

Problem Author:	Kostiantyn Denysov
Prepared by:	Maksym Shvedchenko and Iliia Shevchenko
Editorial by:	Maksym Shvedchenko

Let us note that in order to obtain $M E X$ $i$ we need to have at least $i$ elements in the segment, which means that if we split into $x$ subsegments then their minimal total length will be $\sum_{i = 0}^{x} i = \frac{x ( x + 1 )}{2}$ \sum_{i=0}^{x} i = \frac{x(x+1)}{2} which in turn must be less than $n$ , i.e. $\frac{x ( x + 1 )}{2} < n$ \frac{x(x+1)}{2} < n, hence $x < 2 n$ x < \sqrt{2n}.

Now let us define $d p_{i, j}$ dp_{i,j} — as the number of ways to split the prefix $i$ into exactly $j$ subsegments. Define $L_{i, j}$ L_{i,j} as the leftmost index such that the $M E X$ of the segment $[L_{i, j}, i]$ [L_{i,j},i] equals $j$ , and $R_{i, j}$ R_{i,j} as the rightmost such index. then:

$d p_{i, j} = \sum_{k = L_{i, j} - 1}^{R_{i, j} - 1} d p_{k, j - 1}$ dp_{i,j} = \sum_{k=L_{i,j}-1}^{R_{i,j}-1} dp_{k,j-1}

Now, if we simply search for these boundaries, we get a solution with $O (n^{2} n)$ O(n^2\sqrt n) time and $O (n n)$ O(n\sqrt n) memory. However, note that we can compute the boundaries faster. Observe that $R_{i, j} < R_{i, j - 1}$ R_{i,j} < R_{i,j-1} since the segment must definitely contain all numbers from $1$ to $j - 2$ j-2 and also $j - 1$ j-1. Then let $l a s t_{x}$ last_x be the position of the last occurrence of number $x$ to the left of $i$ . In this case $R_{i, j} = min (R_{i, j - 1}, l a s t_{j - 1})$ R_{i,j} = min(R_{i,j-1},last_{j-1}).

For the left boundary it is even simpler, since the only condition is that the number $j$ must not appear, i.e. $L_{i, j} = l a s t_{j}$ L_{i,j} = last_j.

Now by optimizing the sum computation with prefix sums, we obtain a solution in $n n$ n \sqrt n time and $n n$ n \sqrt n memory. However, we can notice that we may recalculate the DP layer by layer, keeping only the previous layer, which optimizes the memory down to linear.

template<typename it>
struct PrefixSum {
    using T = typename remove_reference<decltype(*declval<it>())>::type;
    vector<T> pref;

    PrefixSum() = default;

    PrefixSum(it first, it last) {
        pref = vector<T>(distance(first, last));
        for (int i = 0; i < len(pref); i++)
            pref[i] = (i == 0 ? 0 : pref[i - 1]) + *(first + i);


    }

    T operator()(int l, int r) {
        assert(l >= 0 && r < len(pref));
        assert(l <= r);
        if (l > r)
            return 0;
        if (l == 0)
            return pref[r];
        return pref[r] - pref[l - 1];
    }
};

signed main() {
    ios_base::sync_with_stdio(false);
    cin.tie(nullptr);
    cout.tie(nullptr);
    int n;
    cin >> n;
    vec<int> a(n);
    cin >> a;
    mint ans = 0;
    int B = sqrt(n*2) + 10;
    vec<mint> cur_dp(n);
    vec<mint> next_dp(n);
    vec<int> max_l(n, 0);
    iota(all(max_l), 0);
    for (int i = 0; i < n; i++) {
        if (a[i] == 0) break;
        cur_dp[i] = 1;
    }
    ans += cur_dp[n - 1];
    PrefixSum ps(all(cur_dp));

    for (int M = 1; M < B; M++) {
        int pos_M = -10;
        int pos_M_plus_1 = -1;
        for (int i = 0; i < n; i++) {
            next_dp[i] = 0;
            if (a[i] == M - 1) {
                pos_M = i;
            } else if (a[i] == M) {
                pos_M_plus_1 = i;
            }
            int R = min(max_l[i], pos_M)-1;
            chmin(max_l[i], pos_M);
            int L = max(0LL, pos_M_plus_1);
            if (R < L) continue;
            next_dp[i] = ps(L, R);
        }
        ans += next_dp[n - 1];
        swap(cur_dp, next_dp);
        ps = PrefixSum(all(cur_dp));
    }
    cout << ans.val() << endl;

}