PROBLEM LINK:
Practice
Contest: Division 1
Contest: Division 2
Contest: Division 3
Contest: Division 4
Author: tanminati
Tester: sushil2006
Editorialist: iceknight1093
DIFFICULTY:
Medium
PREREQUISITES:
Combinatorics
PROBLEM:
For an array B, define f(B) to be the minimum number of times an element of B must be decreased by 1 in order to change the value of \text{MEX}(B).
Given an array A, compute the sum of f(A) across all subsets of A.
EXPLANATION:
Let’s first understand how to compute f(A) for a given array A.
Let M = \text{MEX}(A) initially.
Note that this means A contains at least one copy each of 0, 1, 2, \ldots, M-1, but does not contain M.
We want to change M, which means it must either increase or decrease.
- If we want to increase the MEX, we need to obtain a copy of M in the array.
Since we’re only allowed to decrement elements by 1, clearly the optimal method of doing so is to to choose the smallest existing element \gt M, and repeatedly decrease it till it becomes equal to M.- So, if R is the smallest element larger than M present in A, this requires R - M operations.
- If we want to decrease the MEX, say to x, then we at least need to ensure that there are no copies of x present in the array.
If x \gt 0 then we can turn x into x-1 with a single operation.
If x = 0 then we can’t actually remove any copies of it from A; so it’s impossible to make the MEX 0 if it’s not already 0.- Note that when x \gt 0, we certainly require at least one subtraction operation on every copy of x, for a total of \text{freq}[x] operations.
- Note that these operations are also sufficient, since after them all we’ll have no copies of x so the MEX will surely change to x.
- We don’t care about what the MEX is, only that it changes. So, we can just choose whichever x \lt M is optimal for us.
This leads to a cost of \min(\text{freq}[1], \ldots, \text{freq}[M-1]).
Thus, if M = \text{MEX}(A) and R is the smallest element \gt M present in A, we have
Note that the only situation in which it’s impossible to change the MEX of A is when \min(A) = \max(A) = 0.
For this case alone, we have f(A) = 0.
We now need to sum up this quantity across all subsets of A.
Let c_k denote the number of non-empty subsets of A for which the answer is at least k.
Suppose we’re able to compute all the values of c_k.
Then, observe that the value we’re looking for is exactly
because a subset with answer X will be counted once in each of c_1, c_2, \ldots, c_X; for X times in total.
Let’s now understand how to compute c_k for a fixed k.
We’re looking for subsets whose answer is at least k.
This means the subset must satisfy \min(\text{freq}[1], \ldots, \text{freq}[M-1], R - M) \ge K, where M is its MEX and R is the smallest element present in the subset larger than M.
The min of several values can be \ge K if and only if all of those values are \ge K.
Thus, we require that all of \text{freq}[1], \ldots, \text{freq}[M-1] are \ge K, as well as R-M \ge K \implies R \ge M+K.
This is hard to deal with if we don’t even know M.
So, let’s fix the value of M, the MEX of the subset.
Then, we want subsets such that:
- Each of 1, 2, 3, \ldots M-1 occur at least K times each.
- The MEX of the subset is M (meaning M doesn’t appear, but 0 does appear when M \gt 0).
- The next smallest element after M (if it exists) is at least M+K.
Note that the occurrences of each element are now quite independent of each other, so we can look at each one independently.
Specifically, if we look at value x, with frequency f in A,
- If x = 0, then if M = 0 the subset must not contain it at all (one possibility), and if M \gt 0 then the subset must contain a positive number of occurrences (so 2^f - 1 possibilities.)
- If 1 \le x \lt M, the subset must contain at least K occurrences of it.
- There are \binom{f}{y} ways to choose y occurrences, so there are \binom{f}{K} + \binom{f}{K+1} + \ldots + \binom{f}{f} total ways.
Note that this is a suffix sum of binomial coefficients.
- There are \binom{f}{y} ways to choose y occurrences, so there are \binom{f}{K} + \binom{f}{K+1} + \ldots + \binom{f}{f} total ways.
- If M \le x \lt M+K, then x must not appear at all.
This gives only one possibility. - If x \ge M+K, then x can have any number of occurrences.
This gives 2^f possibilities.
This already gives us a slow algorithm to compute c_k, by iterating over M and then iterating over values of x to compute the appropriate multiplier.
For the given constraints, this needs to be sped up however.
Let’s work with just a fixed k (for c_k) and MEX M for now.
Note that since all values \ge M+K can be chosen freely, we don’t need to look at them individually.
Instead, if s denotes the total number of elements present that are \ge M+K, we simply obtain a multiplier of 2^s since any subset of them can be freely chosen.
(Note that when M \le 1 the multiplier is 2^s - 1 instead, since we need to choose a non-empty subset of these elements given that we can’t actually operate on anything smaller than M.)
Finding s is simple: build suffix sums on the frequency array of A.
That leaves the values in [1, M-1] which we need to deal with.
For each of them, we need to compute some suffix sum of binomial coefficients; which is traditionally a hard problem and can’t be done with some formula in constant time.
Consider instead the following algorithm.
When we’re at k, let’s define an array y as:
to be the appropriate multiplier.
If we know this array, the value we’re looking for once M is fixed will be y_1\cdot y_2 \cdot\ldots\cdot y_{M-1}.
Suppose we know the y_i values for k, and we want to move to k+1 now.
Observe that the array y doesn’t change too much.
In particular, y_i must only change to y_i - \binom{\text{freq}[i]}{k}, i.e. we only subtract out the term corresponding to choosing exactly k elements, which can be done in constant time.
Further, note that we’re only interested in the prefix product of the y_i values, and not each of them individually.
In particular, note that if y_i = 0 for some i, it will remain 0 in the future; and also the values at future indices just don’t matter anymore.
This allows us to use an early-break ‘heuristic’: when moving from k to k+1, update y_i in order of i = 1, 2, 3, \ldots.
If y_i = 0 after the update, immediately break the loop.
This simple-looking heuristic in fact greatly improves our complexity!
Observe that any index i will only be updated \text{freq}[i] times before its y_i value becomes 0; at which point we’ll never update it again.
So, the total number of updates is bounded by
Since each update is just the subtraction of a single binomial coefficient (which can be done in constant time), the algorithm thus runs in linear time overall!
Putting everything together, we have the following algorithm:
- Let y_i denote the number of ways of choosing a subset of the occurrences of i.
Initially, y_i = 2^{\text{freq}[i]} - 1. - Iterate k = 1, 2, 3, \ldots, N
- With k fixed, iterate M = 0, 1, 2, \ldots, N
- While iterating, maintain the prefix product P = y_1\cdot y_2\cdot\ldots y_{M-1}
- The number of subsets with the current (k, M) pair is
P\cdot (2^{\text{freq}[0]}-1) \cdot 2^s
where s denotes the number of elements \ge k+M present in the array. - Note that for M = 0 and M = 1 this formula will vary slightly, since you must ensure a positive number of elements \ge k+M are picked.
- Also note that if P becomes equal to 0, you can break out of the loop immediately - further M will have no contribution.
- After processing all M, we move on to updating the y_i values.
To do this, iterate i = 1, 2, \ldots and update y_i to y_i - \binom{\text{freq}[i]}{k}
If y_i = 0 after the update, break immediately.
As we saw earlier, the total number of updates to y_i is bounded by N as long as early breaks are made.
Note that the same proof also shows that iterating through M till the prefix product becomes 0 and then breaking out results in linear time overall too.
So, the algorithm runs in \mathcal{O}(N) time (or \mathcal{O}(N\log{MOD}) time if you don’t precompute inverses and powers and find them on-the-fly.)
Depending on your implementation, you might also need to specially handle the M = 0 and M = 1 cases in different ways - take care of that appropriately.
TIME COMPLEXITY:
\mathcal{O}(N) per testcase.
CODE:
Editorialist's code (C++)
// #include <bits/allocator.h>
// #pragma GCC optimize("O3,unroll-loops")
// #pragma GCC target("avx2,bmi,bmi2,lzcnt,popcnt")
#include "bits/stdc++.h"
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);
const int mod = 998244353;
auto mpow = [&] (int a, int n) {
int r = 1;
while (n) {
if (n & 1) r = (1ll * r * a) % mod;
a = (1ll * a * a) % mod;
n /= 2;
}
return r;
};
const int N = 1'000'005;
vector<int> fac(N, 1);
for (int i = 1; i < N; ++i) fac[i] = (1ll * fac[i-1] * i) % mod;
auto inv = fac;
inv.back() = mpow(inv.back(), mod-2);
for (int i = N-2; i >= 0; --i) inv[i] = (1ll * (i+1) * inv[i+1]) % mod;
auto C = [&] (int n, int r) {
if (n < r or r < 0) return 0ll;
return ((1ll * fac[n] * inv[r]) % mod * inv[n-r]) % mod;
};
vector p2(N, 1);
for (int i = 1; i < N; ++i) p2[i] = (2 * p2[i-1]) % mod;
int t; cin >> t;
while (t--) {
int n; cin >> n;
vector a(n, 0), freq(n+1, 0);
for (int &x : a) {
cin >> x;
++freq[x];
}
ll ans = 0;
vector val(n+1, 0);
for (int i = 0; i <= n; ++i) {
val[i] = p2[freq[i]] - 1;
}
vector suf = freq;
for (int i = n-2; i >= 0; --i)
suf[i] += suf[i+1];
suf.push_back(0);
for (int k = 1; k <= n; ++k) {
int prod = 1;
for (int i = 0; i <= n; ++i) {
int r = min(n, i+k);
if (i <= 1) ans += 1ll * prod * (p2[suf[r]] - 1);
else ans += 1ll * prod * p2[suf[r]];
ans %= mod;
prod = (1ll * prod * val[i]) % mod;
if (prod == 0) break;
}
for (int i = 1; i < n; ++i) {
val[i] = (val[i] + mod - C(freq[i], k)) % mod;
if (val[i] == 0) break;
}
}
cout << ans % mod << '\n';
}
}