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.
- This is because one balloon from each previous person has hit person i.
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)