Official Analysis (C++, Java, Python)

Explanation

First, observe that the maximum possible amount of milk is generated by assigning the largest multipliers to the cows with the highest milk production values. This is achieved by sorting the array of milk production values in nondecreasing order.

Let's observe the effect of changing an element when maintaining sorted order and how this affects the value of the original value of $T$. Suppose we consider a query that changes $a_i$ to $v$. Let the positions of $a_i$ and $v$ in the sorted version of $a$ be $x$ and $y$, respectively.

If we consider the sorted version of $a$, then there are two cases:

1. If $v \ge a_i$, $a_i$ moves right and the indices $x+1,\ x+2,\ \dots,\ y$ are shifted left.
2. If $v < a_i$, $a_i$ moves left and the indices $y,\ y+1,\ \dots,\ x$ are shifted right.

For the first case, the multipliers of the elements shifted left decrease by 1, so $T$ decreases by the sum of those elements, which can be found using prefix sums. We also have to account for the contributions of the removal of the original element and addition of the new element — we can find their positions using binary search. The second case can be handled similarly.

Implementation

Time Complexity: $\mathcal{O}((N + Q)\log N)$

Copy
import bisect

n = int(input())
a = [int(x) for x in input().split()]

# Sorting and prefix sums
sorted_a = sorted(a)
p_sum = [0]
for i in range(n):
	p_sum.append(p_sum[-1] + sorted_a[i])