PROBLEM LINK:
Practice
Contest: Division 1
Contest: Division 2
Contest: Division 3
Contest: Division 4
Author: raysh07
Tester: iceknight1093
Editorialist: iceknight1093
DIFFICULTY:
TBD
PREREQUISITES:
Prefix sums
PROBLEM:
Alice and Bob play a game on an array, taking turns repeatedly reducing elements by 1 as long as the MEX of the array doesn’t change.
The loser is whoever is unable to make a valid move.
Given an array, count the number of subarrays that would result in Alice winning.
EXPLANATION:
From the easy version, we know the solution to a single array is as follows:
- The total number of moves is fixed irrespective of player choice.
So, the winner is determined purely by the (parity of the) number of moves. - To count the number of moves: if M = \text{MEX}(A) then,
- Every element \gt M gets reduced to M+1.
- Every element \lt M gets reduced to 0, except for one copy each of 0, 1, \ldots, M-1.
We’ll use this information to solve across subarrays.
The first step is to rewrite the cost into something a bit nicer to work with.
With M = \text{MEX}(A) we had:
where f_x denotes the number of occurrences of x in the array.
Expanding this out a bit, and using the fact that f_M = 0, we obtain
Looking at these three terms,
- The first term is simply the sum of the whole array, because x\cdot f_x means “add x once for each time it appears.”
- The second term is 0+1+\ldots + M-1 which equals \frac{M\cdot (M-1)}{2}.
- The third term equals M+1 multiplied by the number of elements in the array greater than M.
Let’s fix a value of M, and try to solve across all subarrays with MEX equal to M.
To do this, we will first do a little setup.
Define a new array B of length N such that:
- If A_i \lt M, then B_i = A_i.
- If A_i \gt M, then B_i = A_i - (M+1).
- If A_i = M then B_i = 0.
B_i represents the contribution of the i-th element to any subarray having mex M; in particular with reference to the sum above this is the element’s contribution to the first and third terms, while the second is a constant anyway.
Now, suppose we fix R to be the right endpoint of the subarray under consideration.
The left endpoints L such that \text{MEX}(A_L, \ldots, A_R) = M form a contiguous range.
Specifically,
- Let x \le R be the largest index such that A_x = M.
- Let y \le R be the largest index such that 0, 1, \ldots, M-1 all appear among elements A_y, \ldots, A_R.
- Then the valid left endpoints L are exactly x \lt L \le y.
- This is because for L \le x the subarray will contain M and so surely can’t have a MEX of M, while if L \gt y then some element smaller than M is missing so again the MEX can’t be M.
Finding these bounds x and y is not too hard either.
To find x, simply store the index of the closest occurrence of M as you sweep over R.
To find y, store the nearest occurrences of all elements \lt M, and take the largest of these nearest occurrences - this can be done in \mathcal{O}(\log N) time using some data structure (like a priority queue), or even in amortized constant time using two pointers, noting that as R increases y can also only increase.
Once x and y are known, the remainder is not too hard.
Recall the array B, which we created to represent the ‘contribution’ of each element.
So, for a fixed left endpoint L, we want to know the parity of
M\cdot (M-1)/2 is a constant, so we only care about the parity of B_L + \ldots + B_R.
This can be rewritten using prefix sums as P_R - P_{L-1}, where P_i = B_1 + \ldots + B_i.
With R fixed, we hence only need to care about the parity of P_{L-1} across all x \lt L \le Y.
This can be found in constant time by building prefix sums over the array of parities of P_i.
Thus, for a fixed value M, we’re able to solve for all subarrays that have mex M in \mathcal{O}(N) or \mathcal{O}(N\log N) time.
We now use the fact that A_i \le 100 to note that the MEX of any subarray is no more than 101; and so simply running this algorithm for each MEX is fast enough.
TIME COMPLEXITY:
\mathcal{O}(N \cdot \max(A)) per testcase.
CODE:
Editorialist's code (C++)
// #pragma GCC optimize("O3,unroll-loops")
// #pragma GCC target("avx2,bmi,bmi2,lzcnt,popcnt")
#include "bits/stdc++.h"
#include <cassert>
using namespace std;
using ll = long long int;
mt19937_64 RNG(chrono::high_resolution_clock::now().time_since_epoch().count());
int main()
{
ios::sync_with_stdio(false); cin.tie(0);
int t; cin >> t;
while (t--) {
int n; cin >> n;
vector a(n, 0);
for (int &x : a) cin >> x;
vector b(n+1, 0), pref(n+1, 0), ppref(n+1, 0), near(105, 0);
ll ans = 0;
for (int m = 0; m < min(n, 100) + 5; ++m) {
for (int i = 0; i < n; ++i) {
if (a[i] < m) b[i+1] = a[i];
else if (a[i] > m) b[i+1] = a[i]-m-1;
else b[i+1] = 0;
pref[i+1] = (pref[i] + b[i+1]) % 2;
ppref[i+1] = ppref[i] + pref[i+1];
}
near.assign(105, -1);
int lo = 1, hi = 1, have = 0;
for (int i = 1; i <= n; ++i) {
have += (a[i-1] < m) and (near[a[i-1]] == -1);
near[a[i-1]] = i;
if (a[i-1] == m) lo = i+1;
if (have < m) continue;
while (hi < i) {
if (a[hi-1] < m and near[a[hi-1]] == hi) break;
++hi;
}
if (lo > hi) continue;
// lo...hi valid as left
// pref[i] - pref[l-1] - m*(m-1)/2 should be odd
int cur = (pref[i] ^ (1ll * m * (m-1)/2) ) % 2;
if (cur == 0) {
// count ones
ans += ppref[hi-1] - ppref[lo-1] + pref[lo-1];
}
else {
// count zeros
ans += hi-lo+1 - (ppref[hi-1] - ppref[lo-1] + pref[lo-1]);
}
}
}
cout << ans << '\n';
}
}