PROBLEM LINK
Author: Hruday Pabbisetty
Tester: Alexey Zayakin
Editorialist: Jakub Safin
DIFFICULTY
SIMPLE
PREREQUISITIES
sorting, greedy algorithms
PROBLEM
You are given a multiset of non-negative integers of size N; maximise its minimum excluded non-negative integer by adding at most K integers to it.
QUICK EXPLANATION
Sort S_1,\dots,S_N, remove duplicate elements. Greedily add numbers between elements S_i, the answer is the smallest number you can’t add; it’s possible to add all numbers between two consecutive elements at once.
EXPLANATION
If we want the mex to be \ge k, we need to have all integers 0 through k-1 in the multiset.
The most straightforward algorithm would be to go through non-negative integers in increasing order. We can ignore integers that are already in the multiset. Whenever encountering an integer x that’s not in the multiset, we need to add it - if we didn’t do it, the mex would be at most x. The first such integer which we can’t add to our multiset (because we already added K of them) is the maximum possible mex.
If we simply mark the numbers from the initial multiset in an array of size \mathrm{max}S, we can determine if a number is in the initial multiset in O(1) and the time complexity of this greedy algorithm is O(\mathrm{max}S+K+N), since we only need to process at most K+N integers before stopping; the memory complexity is O(\mathrm{max}S).
The constraints are low enough that this is enough to solve the problem, but we can do better.
A faster solution
Let’s make this algorithm independent on \mathrm{max}S.
We can sort the array S_{1..N} in non-decreasing order and discard duplicate elements, since the mex doesn’t depend on how many times each number occurs in the multiset, only which numbers; as we go through all non-negative integers, we can remember the index of the last element S_i we encountered (or i=0 if we haven’t encountered any yet). If i < N, then we know that the next number we could encounter that’s in the multiset is S_{i+1}; when that happens, we can just increment i by 1.
With this approach, we only need O(1) time with O(N) memory to check if a number is in the initial multiset and sorting takes O(N\log{N}), so the whole algorithm takes O(N\log{N}+K) time and O(N) memory.
An even faster solution
Let’s make it independent on K, too.
Look at what we’re doing so far after sorting S and removing duplicates: add all numbers from 0 up to S_1-1 if we can, skip S_1, add all numbers from S_1+1 up to S_2-1 if we can, skip S_2, etc. Instead of adding numbers in a range [x,y] one by one, why not just subtract y-x+1 from K, since it has the same effect?
So, we can simply go through the numbers S_{1..N} and for each S_i, subtract S_i-1-S_{i-1} from K (we can consider S_0=-1 so that we’re subtracting S_1 for i=1). As soon as we can’t do this, because K < S_i-1-S_{i-1}, the only numbers we can add are S_{i-1}+1 through S_{i-1}+K and the answer will be S_{i-1}+K+1.
Naturally, if this didn’t happen for any i, then we should be adding numbers starting from S_N+1; the first number we couldn’t add is S_N+1+K, so that will be the answer.
This solution takes just O(N\log{N}) time and O(N) memory. There’s an additional \log factor, but we can solve the problem even for very large values of K and S_{1..N}.