PROBLEM LINK:
Practice
Contest: Division 1
Contest: Division 2
Contest: Division 3
Contest: Division 4
Author: iceknight1093
Tester: sushil2006
Editorialist: iceknight1093
DIFFICULTY:
Easy
PREREQUISITES:
Greedy
PROBLEM:
You have an array A.
You can do the following at most once:
- Choose at most K indices.
Add 1 to the elements at the chosen indices.
Maximize \text{mode}(A) after the operation.
EXPLANATION:
Let \text{freq}[x] denote the number of times x appears in A.
A value M is a mode of A if and only if \text{freq}[M] \ge \text{freq}[x] for all x.
Since we can only increment elements by 1, incrementing x is equivalent to being able to “transfer” value from \text{freq}[x] to \text{freq}[x+1].
We want to maximize \text{mode}(A).
Suppose we fix M to be the candidate for the mode.
Let’s further fix the desired final frequency for M, say F.
Note that fixing F also uniquely determines the number of moves of the form M-1 \to M we make.
Now that F is fixed, we need to ensure that M is actually the mode.
This means we need to ensure that every other element has a frequency that’s no more than F.
The only way to decrease an element’s frequency is to increment the element itself by 1.
So,
- If \text{freq}[1] \gt F, then we need to increment at least (F - \text{freq}[1]) copies of 1 to 2 to make it satisfy the frequency condition.
This will change the frequency of 2.- If \text{freq}[1] \le F instead, we don’t need to touch the 1's at all; since our moves are limited this is fine.
- Next, if \text{freq}[2] \gt F, similarly we need to increment (F - \text{freq}[2]) copies of 2 to become 3.
However, note that we can only increment previously existing copies of 2, i.e. at most \text{original\_frequency}[2] copies.
Note that if this means we don’t have enough copies of 2 to increment, it’s impossible to have \le F copies of 2, so M cannot be made the mode (while having frequency F). - Similarly, for each of 3, 4, 5, \ldots, M-2, if their current frequencies (after accounting for increases from the previous values) exceed F, we’ll need to increment several copies of the element; but not more than the number of initial copies.
- Note that we don’t check for the value M-1, since its number of increments is already determined by F itself.
So, we only need to check if the final frequency of M-1 (after increments to M-2) remains \le F.
- Note that we don’t check for the value M-1, since its number of increments is already determined by F itself.
This simple greedy algorithm gives us a lower bound on the number of elements that need to be incremented, since we worked on exactly everything that we needed to.
Note that we don’t need to perform the check with values \gt M.
This is because, if one of these values has a higher frequency, M will not be a mode - but then certainly something larger than M will be a mode; which is good for us since our aim is to maximize a mode.
Further, if this case happens, it will be considered when processing said larger value anyway, so we don’t lose any information by just ignoring values \gt M.
Note that the above algorithm gives us a simple \mathcal{O}(N) greedy check for whether M (or rather, a value \ge M) can be a mode, with M having frequency F.
Simply running this for all pairs of (M, F) leads to a solution with a complexity of \mathcal{O}(N^3) on the surface.
However, we can in fact prove stricter bounds on the complexity!
In particular, note that once M is fixed, there are only \text{freq}[M-1] + 1 possible values for F, since we’re bounded by the number of elements M-1 that we increment to M.
This means the number of pairs (M, F) we need to check is really
because the sum of all frequencies is just N.
Since each check is performed in linear time, our algorithm actually runs in \mathcal{O}(N^2) which is now fast enough for the given constraints.
It’s possible to make one further observation: if M is fixed, it’s optimal to first convert as many elements as possible to M before doing anything else (this is not hard to prove.)
This means we only need to check for F = \text{freq}[M] + \min(K, \text{freq}[M-1]).
However, there are still N choices of M, so the complexity doesn’t actually change upon using this observation.
TIME COMPLEXITY:
\mathcal{O}(N^2) per testcase.
CODE:
Editorialist's code (PyPy3)
import sys
input = sys.stdin.readline
for _ in range(int(input())):
n, k = map(int, input().split())
a = list(map(int, input().split()))
freq = [0]*(2*n+5)
for x in a: freq[x] += 1
ans = 0
for x in range(1, 2*n+2):
for ops in range(freq[x-1]+1):
curf = freq[x] + ops
if curf == 0: continue
# keep everything else <= curf
add, good, req = 0, 1, 0
for y in range(x-1):
send = max(0, freq[y] + add - curf)
req += send
good &= send <= freq[y]
add = send
good &= freq[x-1] - ops + add <= curf
if good and req+ops <= k: ans = x
print(ans)