PROBLEM LINK:
Practice
Contest: Division 1
Contest: Division 2
Contest: Division 3
Contest: Division 4
Author: notsoloud
Tester: mexomerf
Editorialist: iceknight1093
DIFFICULTY:
cakewalk
PREREQUISITES:
None
PROBLEM:
Returning a book to a library incurs a penalty equal to the number of days late it was returned.
However, if multiple books are returned on the same day, only the maximum penalty among them needs to be paid.
Given N books, the i-th of which was returned i days late and was returned on day A_i, find the total penalty that must be paid.
EXPLANATION:
Let \text{max\_penalty}[x] denote the maximum penalty incurred by a book that’s returned on day x.
If we can compute this for all x, the answer is simply the sum of the \text{max\_penalty} array.
Computing this array can be done as follows:
- Initialize all its elements to 0.
- Then, for each i from 1 to N in order, set \text{max\_penalty}[A_i] to i.
Finally, compute the sum of \text{max\_penalty}.
This can be implemented in \mathcal{O}(N\log N) time using map to store the array, or \mathcal{O}(N + 10^5) time by creating an array of length 10^5 for each test case (10^5 is needed because that’s the maximum possible value in A), which is fast enough since there are only 1000 tests.
TIME COMPLEXITY:
\mathcal{O}(N\log N) or \mathcal{O}(N) per testcase.
CODE:
Editorialist's code (PyPy3)
for _ in range(int(input())):
n = int(input())
a = list(map(int, input().split()))
ans = 0
max_penalty = [0]*(max(a) + 1)
for i in range(1, n+1):
max_penalty[a[i-1]] = i
print(sum(max_penalty))