PROBLEM LINK:
Practice
Contest: Division 1
Contest: Division 2
Contest: Division 3
Contest: Division 4
Author: iceknight1093
Tester: raysh07
Editorialist: iceknight1093
DIFFICULTY:
Medium
PREREQUISITES:
Sorting, binary search
PROBLEM:
Given an array B of length N, define f(B) as follows:
- Start with an array A of 2N zeros.
- For each i = 1, 2, \ldots, N in order, choose an index j such that A_j \le B_i and increment A_j by 1.
- f(B) denotes the maximum possible value of \text{MEX}(A) after all N increments have been made; assuming they were chosen optimally.
Given K, find any rearrangement R of B such that f(R) = K.
EXPLANATION:
First, let’s attempt to compute f(B) for a fixed B.
If we want \text{MEX}(A) = K, then we need to form each of the values 0, 1, 2, \ldots, K-1 in A.
(Technically we need to also ensure that K doesn’t appear; but since our objective is to maximize the MEX it’s not an issue even if it does.)
Since A starts with 2N zeros, and we have only N increments, A will definitely contain some zeros in the end. So, we only need to worry about the positive elements.
To keep things simple, we’ll try to make A_i = i-1 for each 1 \le i \le K.
Suppose we’re looking at B_i (having already processed B_1, \ldots, B_{i-1}).
We now need to choose some index j such that A_j \le B_i and increment A_j by 1.
Which j do we choose?
First, we can clearly ignore all j such that A_j = j-1; since they’ve reached their “destination” values already.
Among the remaining choices of j, it’s optimal to choose whichever one has the largest value of A_j (if there are ties, choose the largest j \le K as well.)
That is, we have a simple greedy algorithm: always increment the largest value that we’re able to increment (and that still needs incrementing.)
The optimality of this choice can be proven using an exchange argument - if we don’t use B_i here, then we’re using it for something smaller; but this increment needs to be done anyway so a later element must do it. However in that case we could use B_i here and the later element for the smaller value without losing anything.
(Note that it’s possible no valid j exists, for example when B_i is too small late into the array. In this case we can essentially just ignore B_i and not increment anything; by pretending we’re incrementing some random 0 after index K.)
Finding the optimal element to increment can be done quickly using binary search, because it can be verified that our greedy algorithm always maintains A_1 \le A_2 \le\ldots\le A_K.
So, checking for whether f(B) \ge K can be done in \mathcal{O}(N\log K) time.
This allows us to use binary search to compute the actual value of f(B) in \mathcal{O}(N\log ^ 2 N ) time.
We now know how to compute f(R) for a fixed rearrangement R.
Let’s think of the extremes: what are the lowest and highest possible values of f(R) across all rearrangements?
Since it’s clear that larger elements are “more useful”, intuitively it makes sense to keep them for the end and “use up” smaller elements first.
Ergo, the natural guess is that:
- f(R) is minimized when R is sorted in descending order.
- f(R) is maximized when R is sorted in ascending order.
This is indeed true - and can be proved by noting that if R_i \gt R_{i+1} then swapping them can only increase the answer since at the very least the same set of increments can be made.
Let mn denote the minimum possible value of f(R) across all rearrangements, and mx denote the maximum possible value.
Each one can be computed in \mathcal{O}(N\log ^2 N) time using the algorithm above; since we know which array attains which.
Now, clearly if K \lt mn or K \gt mx no rearrangement can ever attain a score of K.
Thus, we only need to worry about mn \le K \le mx.
Luckily, it turns out that all of these are achievable!
One construction is as follows.
Suppose we have some rearrangement R.
Consider what happens if we take any one single element and move it to some arbitrary position of the array, to obtain the new array R_2.
Then, f(R_2) \ge f(R)-1, because at worst we’ll lose the contribution of the moved element which can affect the creation of only one of the values among 0, 1, \ldots, f(R)-1.
Further, if it does affect a value, we can safely assume it affects the creation of f(R)-1 by moving operations around a bit.
Thus we’re still able to create 0, 1, \ldots, f(R)-2 at least, which tells us that f(R_2)\ge f(R)-1.
Now, let’s start with R being in ascending order, giving f(R) = mx.
If we repeatedly take the last element of R and move it to the front of the array (but placed after previously moved elements), after N steps we’ll end up with f(R) = mn since we’ll reach descending order.
For example, if we start with [1, 2, 3, 4, 5], then we sequentially make it
- [1, 2, 3, 4, 5]
- [5, 1, 2, 3, 4]
- [5, 4, 1, 2, 3]
- [5, 4, 3, 1, 2]
- [5, 4, 3, 2, 1]
However, recall that the maximum attainable MEX changes by 1 at each step.
So, moving from mx to mn, we’ll definitely hit everything inbetween!
To find where it becomes exactly K, we can binary search because the transformation is monotonic.
Note that if implemented directly, this gives a complexity of \mathcal{O}(N\log ^3 N), because we have one log from the binary search and then a check in \mathcal{O}(N\log^2 N) each time.
However, this can be optimized a bit.
In particular, we don’t need a \mathcal{O}(N\log^2 N) check each time.
Or rather, we don’t care about the exact score of the current arrangement - just whether the score is \ge K or not.
However, we already know how to perform that check in \mathcal{O}(N\log N) time; as the first thing we discussed here.
Thus, we can skip one binary search to attain a complexity of \mathcal{O}(N\log N\log K), which is quite fast.
In fact, because we need 1+2+3+\ldots + K-1 increments to attain a score of K, any valid K must be bounded by about \sqrt{2N}.
So, the complexity is \mathcal{O}(N\log N\log (\sqrt N)), which is technically the same as \mathcal{O}(N\log^2 N) but with better constant factor; and runs rather quickly in practice.
TIME COMPLEXITY:
\mathcal{O}(N\log N \log(\sqrt N)) 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());
void Solve()
{
int n, k; cin >> n >> k;
if (k > 800){
cout << -1 << "\n";
return;
}
vector <int> b(n);
for (auto &x : b) cin >> x;
auto check = [&](vector <int> a, int k){
if (k <= 1) return true;
vector <int> f(k);
for (int i = 1; i < k; i++) f[i] = 0;
set <int> s;
for (int i = 1; i < k; i++) s.insert(i);
for (int i = 0; i < n; i++){
int x = a[i];
int lo = 0, hi = k - 1;
while (lo != hi){
int mid = (lo + hi + 1) / 2;
if (f[mid] <= x) lo = mid;
else hi = mid - 1;
}
int pos = lo;
auto id = s.upper_bound(pos);
if (id != s.begin()){
--id;
lo = *id;
f[lo]++;
if (f[lo] == lo) s.erase(lo);
}
}
for (int i = 0; i < k; i++) if (f[i] < i){
return false;
}
return true;
};
sort(b.begin(), b.end());
if (!check(b, k)){
cout << -1 << "\n";
return;
}
reverse(b.begin(), b.end());
if (check(b, k + 1)){
cout << -1 << "\n";
return;
}
reverse(b.begin(), b.end());
int lo = 0, hi = n;
while (lo != hi){
int mid = (lo + hi + 1) / 2;
vector <int> a;
for (int i = 0; i < mid; i++) a.push_back(b[n - 1 - i]);
for (int i = 0; i < n - mid; i++) a.push_back(b[i]);
if (check(a, k)){
lo = mid;
} else {
hi = mid - 1;
}
}
for (int i = 0; i < lo; i++) cout << b[n - 1 - i] << " ";
for (int i = 0; i < n - lo; i++) cout << b[i] << " ";
cout << "\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;
}