MAKEDISTK - Editorial

PROBLEM LINK:

Practice
Contest: Division 1
Contest: Division 2
Contest: Division 3
Contest: Division 4

Author: iceknight1093
Tester: sushil2006
Editorialist: iceknight1093

DIFFICULTY:

Simple

PREREQUISITES:

Sorting

PROBLEM:

You’re given an array A of length N, and K.
You can choose upto K indices and increment the elements at each of their positions by 1.
Find the minimum number of operations needed to ensure that the elements of A are pairwise distinct.

EXPLANATION:

Let’s sort the array A in non-decreasing order first. This obviously does not change the answer.
From now on, we assume A_1 \le A_2 \le\ldots\le A_N.

Let B_i denote the final value taken by the i-th element.
We want all the B_i values to be distinct.
It’s not hard to see that it’s optimal to just have B_1 \lt B_2 \lt\ldots\lt B_N.

Why?

Suppose we have B_i \gt B_{i+1}.
Then, because A_i \le A_{i+1}, there are surely at least (B_i - B_{i+1}) operations where we incremented the i-th element but not the (i+1)-th.

Simply choose any (B_i - B_{i+1}) of these operations, and increment the (i+1)-th element instead of the i-th one in them.
This will effectively swap the values of B_i and B_{i+1} without changing the number of operations.

Doing this repeatedly will allow us to reach a sorted array without increasing the number of operations.

With this knowledge, let’s try to just find the optimal final values B_i.


B_1 is obvious: we simply have B_1 = A_1, since we can just not increment A_1 ever.
This is allowed because we are allowed to choose at most K elements; so any move that increments A_1 we can simply drop A_1 out of.

Next, we look at B_2.
B_2 must definitely be larger than B_1, so B_2 \ge B_1 + 1.
However, we must also have B_2 \ge A_2, since we’re incrementing A_2 to reach it.
Thus, B_2 \ge \max(B_1 + 1, A_2).

Now observe that we can simply choose B_2 = \max(B_1+1, A_2) and simply stop incrementing the second element once it reaches this point (again, allowed because at most K elements.)

In fact, applying the same logic tells us that in general, for i \ge 2, we have B_i = \max(B_{i-1}+1, A_i) as the optimal final value for the i-th element.

Thus, we can compute all the final values B_1, \ldots, B_N in linear time by processing from left to right.
Note that these final values are completely independent of K.


Now that the final values are known, let’s compute the number of moves needed.

Define d_i = B_i - A_i to be the number of times the i-th element must be increased.

We now have two restrictions:

  1. Each element can be incremented at most once by a single move.
  2. Each move can increment at most K different elements.

Each of these restrictions gives us a lower bound on the number of moves needed.

The first restriction tells us that we surely need at least d_i operations to make A_i reach B_i.
So, a lower bound on the number of moves is given by

\max(d_1, d_2, \ldots, d_N)

Let M = \max(d_1, \ldots, d_N).

The second restriction tells us that the total number of increments needed can be reduced by at most K in a single move.
Thus, if we let S = d_1 + \ldots + d_N denote the total number of increments needed, S can be reduced by at most K in one move.
So, we need at least

\text{ceil}\left(\frac{S}{K}\right)

moves to process all increments.

We now have two lower bounds on the number of moves needed, so we certainly need to satisfy whichever of them is stricter, i.e. the answer is at least

\max\left(M, \text{ceil}\left(\frac{S}{K}\right)\right)

It can be proved that this number of operations is not just necessary, but also sufficient.

Proof

A simple construction is to always choose the largest K values of d_i and operate on them, hence decrementing them by 1.
(If there are less than K positive values of d_i, just choose all positive ones.)

If there are less than K positive d_i, what we’re doing is clearly optimal since each of them will move closer to the target simultaneously; and so we’ll end up with just \max(d_i) operations.

If there are at least K positive d_i, observe that S decreases by K, so \text{ceil}\left(\frac{S}{K}\right) decreases by 1 for sure.
So, we only need to analyze whether M changes or not.
To do that, we look at two cases.

Case 1: M \lt \text{ceil}\left(\frac{S}{K}\right)
Here, it doesn’t actually matter if M decreases, because the value of \max(M, \text{ceil}\left(\frac{S}{K}\right)) will definitely decrease either way.
So this case is fine - in fact we can just choose any K elements with positive d_i to operate on and get the same result.

Case 2: M \ge \text{ceil}\left(\frac{S}{K}\right)
Here, if we have at most K occurrences of M, our “greedily choose largest K differences” will surely decrement all of them and M will decrease by 1, so we’re done.

What if we have \gt K occurrences of M?
As it turns out, that case is not possible!
That is because M \ge \text{ceil}\left(\frac{S}{K}\right) \ge \frac{S}{K}, so we have M\cdot K \ge S.

If we had more than K occurrences of M, the overall sum would be at least M\cdot (K+1) which is strictly larger than S, contradicting S being the sum of all d_i.
(Note that this only true because all d_i are non-negative.)

This proves that the \max\left(M, \text{ceil}\left(\frac{S}{K}\right)\right) bound on number of moves is indeed attainable.

TIME COMPLEXITY:

\mathcal{O}(N\log N) per testcase.

CODE:

Editorialist's code (PyPy3)
for _ in range(int(input())):
    n, k = map(int, input().split())
    a = list(map(int, input().split()))
    
    a.sort()
    b = a[:]
    for i in range(1, n):
        b[i] = max(b[i-1] + 1, a[i])
    
    mx, sm = 0, 0
    for i in range(n):
        mx = max(mx, b[i] - a[i])
        sm += b[i] - a[i]
    
    print(max(mx, (sm + k - 1) // k))
    
1 Like