PROBLEM LINK:
Practice
Contest: Division 1
Contest: Division 2
Contest: Division 3
Contest: Division 4
Author: iceknight1093
Tester: raysh07
Editorialist: iceknight1093
DIFFICULTY:
Medium
PREREQUISITES:
Combinatorics
PROBLEM:
For an array A, define f(A) as the maximum integer k such that A can be partitioned into k non-empty sets that all have distinct MEX-es.
You’re given an array A. Solve the following problem for each 1 \le i \le N:
- Let M = f(A[1\ldots i])
Count the number of non-empty subsequences S of A[1\ldots i] that satisfy f(S) = M.
EXPLANATION:
Let’s understand how to compute f(A) for a fixed A.
Observe that if we can have k distinct MEX-es, it’s (almost) always possible to have them be the values 0, 1, \ldots, k-1.
This is because:
- If we have some subset S_i with positive MEX but not 0, just take all positive elements of S_i into their own set, and place its occurrences of 0 in whichever set has MEX equal to 1 (or make them a set of their own, if 1 doesn’t exist as a MEX).
This doesn’t reduce the number of distinct MEX-es, but now 0 is included as one of them. - If we have x and y but not x+1 (where x+1 \lt y) as MEXes, it’s easy enough to modify the set with MEX equal to y to make it x+1 instead - move all occurrences of x+1 to the set with MEX 0.
There is only one exception to this: if we have no positive elements, then the only MEX we can create is 1, and 0 cannot be created.
With this in mind, let’s try to find the largest integer k such that all of 0, 1, 2, \ldots, k-1 can be created as MEX-es (the special case of all zeros is trivially handled).
This is easy to analyze:
- We need an occurrence of 0 for each of the sets with MEX 1, 2, \ldots, k-1, so that’s k-1 occurrences in total.
- We need an occurrence of 1 for each of the sets with MEX 2, \ldots, k-1, so that’s k-2 occurrences in total.
\vdots - For any 0 \le i \lt k-1, we need an occurrence of i for each of the sets with MEX i, i+1, \ldots, k-1, so that’s k-1-i occurrences in total.
- Further, for a MEX of 0 we also need an extra positive element, other than the ones used already.
This gives us the following criteria:
- \text{freq}[x] \ge k-1-x for each 0 \le x \lt k-1
- \text{freq}[1] + \text{freq}[2] + \ldots \gt (1 + 2 + 3 + \ldots + k-2)
f(A) then equals the largest k for which this condition is true.
Now, suppose we know f(A) = k. Let’s compute the number of its subsequences that attain said maximum.
The above two conditions in fact tell us exactly what to count:
- For each 0 \le x \lt k-1, we need to pick at least k-1-x occurrences of x.
If there are y occurrences of x, the number of ways of doing this is
\binom{y}{M-1-x} + \binom{y}{M-x} + \ldots + \binom{y}{y} - These counts for each x can all be multiplied together, since they’re independent choices.
- For elements \ge k-1, we can freely pick them in the subsequence, since they don’t affect potential MEX-es.
If there are c elements \ge k-1, that gives us an extra 2^c multiplier.
However, we also need to deal with the condition of \text{freq}[1] + \text{freq}[2] + \ldots \gt (1 + 2 + 3 + \ldots + k-2), i.e. ensure there’s at least one extra non-zero element.
To account for this, note the “bad” case occurs when we pick exactly k-1-x occurrences of x for every x in [1, k-2], as well as picking no items \ge k-1 (and also pick \ge k-1 occurrences of 0).
Counting the number of such configurations is simple enough, it’s again just a product of several binomial coefficients (in the case of 0 alone, it’s a sum of coefficients instead).
We now know how to solve the problem for a fixed array in linear time.
Next, we need to solve for each prefix.
Let’s see what changes as we move from one prefix to another.
For convenience, let’s store M_x as the multiplier corresponding to x (recall that we obtained an independent multiplier for each element \lt k-1).
We also store the product P of all M_x, as well as 2^c, the multiplicative factor.
This means the initial answer is P\cdot 2^c, after which we subtract out the product of some binomial coefficients (which can also be stored).
First, suppose f(A[1\ldots i]) = f(A[1\ldots i+1]) = k, i.e. the maximum doesn’t change.
Then, if we have x = A_{i+1},
- If x \ge k-1, the only change we need to do is make 2^c into 2^{c+1}.
- If x \lt k-1, \text{freq}[x] will increase by 1.
This will change the multiplier corresponding to only x, nothing else changes.
So, what we can do is divide M_x out of P, compute the new value of M_x, and then multiply it back in.
How do we compute the new value of M_x?
That can be done with a bit of combinatorics - essentially, we have an expression of the form
and it turns into
This transformation can be computed in constant time by applying Pascal’s identity of \binom{n}{r} = \binom{n-1}{r} + \binom{n-1}{r-1}, the details are left to you as an exercise.
Overall, in \mathcal{O}(\log{MOD}) time (due to the division), it’s possible to update the values necessary to know the answer, so we’re done.
That leaves the case of f(A[1\ldots i]) \lt f(A[1\ldots i+1])
For convenience, let’s assume f(A[1\ldots i]) = k and f(A[1\ldots i+1]) = k+1.
Looking at what changes:
- For each 0 \le x \lt k-1, the multiplier M_x loses one term of its summation.
This can be updated in \mathcal{O}(1) time for each x, for \mathcal{O}(k) overall. - We obtain a new multiplier M_{k-1}, equal to 2^{\text{freq}[k-1]} - 1 since the subset just needs to include at least one copy of k-1.
This can be computed in logarithmic time (or even constant, with some precomputation). - 2^c will change to 2^{c - \text{freq}[k-1]}, again easy to update.
So, everything can be kept updated in \mathcal{O}(k \log{MOD}) time.
After this, we can then update things with the value of A_{i+1} (which now takes \mathcal{O}(\log{MOD}) time).
If moving from f(A[1\ldots i]) to f(A[1\ldots i+1]) requires an increase of more than 1 in the answer, just perform each increment one at a time till the new value is reached.
Let’s analyze the complexity now.
We do \mathcal{O}(k\log{MOD}) work only when moving from k to k+1, and do logarithmic work at any other time.
So, if M = f(A) is the maximum possible answer, the work we do is \mathcal{O}((N + M^2)\log{MOD}).
However, note that M cannot actually be too large: indeed, to obtain 0, 1, 2, \ldots, M-1 as MEX-es, we need around \mathcal{O}(M^2) elements in the first place.
This means \mathcal{O}(M^2) is bounded by \mathcal{O}(N), meaning our complexity is really just \mathcal{O}(N\log{MOD}), easily fast enough!
Finally, the one missing link is that we need to be able to check if f(A[1\ldots i]) and f(A[1\ldots i+1]) are the same or not, to identify when an increase occurs.
This can be done reasonably easily with a bit of bookkeeping - only \mathcal{O}(k) frequencies matter, along with a suffix sum of frequencies, so you can for example store the count of indices which satisfy the frequency condition for k+1, and update that in \mathcal{O}(1) when moving to i+1.
This needs to be reset when moving from k to k+1, but that can be done in \mathcal{O}(k) so it ends up being \mathcal{O}(N) work overall anyway for the same reason as above.
One thing you might notice is that our solution involves division - specifically, we divide by a certain sum of binomial coefficients, and if this sum ever becomes 0 we’ll land up in trouble.
Luckily, this turns out to not cause any problems.
This is because the sum we divide by will be of the form
for some n \le 5\cdot 10^5 and r \le 1000.
It can be verified locally (there are about 5\cdot 10^5\cdot 1000 possibilities to try, which can be done in a few seconds) that all of these values are never 0 under the modulo, so we’re safe.
TIME COMPLEXITY:
\mathcal{O}(N\log{MOD}) per testcase.
CODE:
Tester's code (C++)
#include <bits/stdc++.h>
using namespace std;
#define int long long
#define INF (int)1e18
mt19937_64 RNG(chrono::steady_clock::now().time_since_epoch().count());
const int mod = 998244353;
struct mint{
int x;
mint (){ x = 0;}
mint (int32_t xx){ x = xx % mod; if (x < 0) x += mod;}
mint (long long xx){ x = xx % mod; if (x < 0) x += mod;}
int val(){
return x;
}
mint &operator++(){
x++;
if (x == mod) x = 0;
return *this;
}
mint &operator--(){
if (x == 0) x = mod;
x--;
return *this;
}
mint operator++(int32_t){
mint result = *this;
++*this;
return result;
}
mint operator--(int32_t){
mint result = *this;
--*this;
return result;
}
mint& operator+=(const mint &b){
x += b.x;
if (x >= mod) x -= mod;
return *this;
}
mint& operator-=(const mint &b){
x -= b.x;
if (x < 0) x += mod;
return *this;
}
mint& operator*=(const mint &b){
long long z = x;
z *= b.x;
z %= mod;
x = (int)z;
return *this;
}
mint operator+() const {
return *this;
}
mint operator-() const {
return mint() - *this;
}
mint operator/=(const mint &b){
return *this = *this * b.inv();
}
mint power(long long n) const {
mint ok = *this, r = 1;
while (n){
if (n & 1){
r *= ok;
}
ok *= ok;
n >>= 1;
}
return r;
}
mint inv() const {
return power(mod - 2);
}
friend mint operator+(const mint& a, const mint& b){ return mint(a) += b;}
friend mint operator-(const mint& a, const mint& b){ return mint(a) -= b;}
friend mint operator*(const mint& a, const mint& b){ return mint(a) *= b;}
friend mint operator/(const mint& a, const mint& b){ return mint(a) /= b;}
friend bool operator==(const mint& a, const mint& b){ return a.x == b.x;}
friend bool operator!=(const mint& a, const mint& b){ return a.x != b.x;}
mint power(mint a, long long n){
return a.power(n);
}
friend ostream &operator<<(ostream &os, const mint &m) {
os << m.x;
return os;
}
explicit operator bool() const {
return x != 0;
}
};
// Remember to check MOD
struct factorials{
int n;
vector <mint> ff, iff;
factorials(int nn){
n = nn;
ff.resize(n + 1);
iff.resize(n + 1);
ff[0] = 1;
for (int i = 1; i <= n; i++){
ff[i] = ff[i - 1] * i;
}
iff[n] = ff[n].inv();
for (int i = n - 1; i >= 0; i--){
iff[i] = iff[i + 1] * (i + 1);
}
}
mint C(int n, int r){
if (n == r) return mint(1);
if (n < 0 || r < 0 || r > n) return mint(0);
return ff[n] * iff[r] * iff[n - r];
}
mint invC(int n, int r){
if (n == r) return mint(1);
if (n < 0 || r < 0 || r > n) return mint(0);
return iff[n] * ff[r] * ff[n - r];
}
mint P(int n, int r){
if (n < 0 || r < 0 || r > n) return mint(0);
return ff[n] * iff[n - r];
}
mint solutions(int n, int r){
// Solutions to x1 + x2 + ... + xn = r, xi >= 0
return C(n + r - 1, n - 1);
}
mint catalan(int n){
return ff[2 * n] * iff[n] * iff[n + 1];
}
};
const int PRECOMP = 3e6 + 69;
factorials F(PRECOMP);
// REMEMBER To check MOD and PRECOMP
void Solve()
{
int n; cin >> n;
vector <int> g(2 * n + 1);
for (int i = 0; i < n; i++){
g[i]++;
}
vector <mint> c(n, 0);
vector <int> freq(n), cut(n);
mint ans = 1, curr = 1;
int p = 0;
vector <mint> p2(n + 1);
p2[0] = 1;
for (int i = 1; i <= n; i++){
p2[i] = p2[i - 1] * 2;
}
auto upd1 = [&](int x){
ans /= p2[freq[x]] - c[x];
curr /= F.C(freq[x], cut[x]);
curr *= F.C(freq[x] + 1, cut[x]);
if (cut[x]){
mint val = c[x];
val *= 2;
val -= F.C(freq[x], cut[x] - 1);
c[x] = val;
}
ans *= p2[freq[x] + 1] - c[x];
};
auto upd2 = [&](int x){
ans /= p2[freq[x]] - c[x];
curr /= F.C(freq[x], cut[x]);
curr *= F.C(freq[x], cut[x] + 1);
mint val = c[x];
val += F.C(freq[x], cut[x]);
c[x] = val;
ans *= p2[freq[x]] - c[x];
};
bool can = false;
for (int i = 1; i <= n; i++){
int x; cin >> x;
can |= x != 0;
upd1(x);
g[freq[x] + x]--;
freq[x]++;
g[freq[x] + x]++;
while (true){
if (g[p] > 0) break;
int extra = freq[0] - p - 1;
int ac = i - (p + 1) * (p + 2) / 2;
if (ac == extra && p > 0){
break;
}
for (int i = 0; i <= p; i++){
upd2(i);
cut[i]++;
}
p++;
}
mint res = ans;
if (can){
mint ok = curr;
ok /= F.C(freq[0], cut[0]);
ok *= p2[freq[0]] - c[0];
res -= ok;
}
cout << res << " \n"[i == n];
}
}
int32_t main()
{
auto begin = std::chrono::high_resolution_clock::now();
ios_base::sync_with_stdio(0);
cin.tie(0);
int t = 1;
// freopen("in", "r", stdin);
// freopen("out", "w", stdout);
cin >> t;
for(int i = 1; i <= t; i++)
{
//cout << "Case #" << i << ": ";
Solve();
}
auto end = std::chrono::high_resolution_clock::now();
auto elapsed = std::chrono::duration_cast<std::chrono::nanoseconds>(end - begin);
cerr << "Time measured: " << elapsed.count() * 1e-9 << " seconds.\n";
return 0;
}