BOP3E - Editorial

PROBLEM LINK:

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

Author: kingmessi
Tester: pols_agyi_pols
Editorialist: iceknight1093

DIFFICULTY:

Cakewalk

PREREQUISITES:

None

PROBLEM:

N friends play a game of balloon smash.

For each i = 1, 2, 3, \ldots, N in order, friend i (if they’re still at the party) will hit every other remaining friend with one balloon, and then leave the party.

For each j, if friend j gets hit by A_j water balloons, they will immediately leave the party and cannot be hit further.
Find the number of balloons each friend will be hit by in the end.

EXPLANATION:

Let’s think of a slow solution first.

We define an array H, where H_i denotes the number of balloons that the i-th friend has been hit by so far.
Initially, H_i = 0 for all i.

Now, for each i = 1, 2, 3, \ldots, N in order,

  • If H_i = A_i, this friend has been hit by enough balloons that they have left the party already.
    So, nothing happens.
  • If H_i \lt A_i, this friend is still at the party.
    They will hit every remaining friend with a balloon.
    • So, for each j such that j > i and H_j \lt A_j, the value of H_j will increase by 1.

This can easily be implemented to run in \mathcal{O}(N^2) time, with the final answer being the array H.

To improve the complexity, we make the following observation:

  • When we’re looking at friend i, let K denote the number of friends before this who threw balloons.
    Then, at this instant we have H_i = \min(A_i, K).
    • This is because one balloon from each previous person has hit person i.
      If this count is A_i or larger, friend i would’ve left the party as soon as the A_i-th balloon hit.
      Otherwise, they’ve been hit by every balloon so far.

So, rather than updating the entire array H every time, we can build it dynamically while iterating through friends.

That is, let K denote the number of friends who have thrown balloons so far.
Then, for each i = 1, 2, \ldots, N,

  • If K \ge A_i, person i has been hit by A_i balloons and left the party.
    Their final H_i value is A_i itself, so we just output this.
    K doesn’t change.
  • If K \lt A_i, this person has been hit by K balloons and is still at the party.
    Their final H_i value is K, so we output K.
    Then K is incremented by 1, since they throw a balloon at everyone else.

This solves the problem in linear time.

TIME COMPLEXITY:

\mathcal{O}(N) per testcase.

CODE:

Editorialist's code (PyPy3)
for _ in range(int(input())):
    n = int(input())
    a = list(map(int, input().split()))
    
    k = 0
    ans = [0]*n
    for i in range(n):
        if a[i] <= k:
            ans[i] = a[i]
        else:
            ans[i] = k
            k += 1
    print(*ans)