Accepted Code — 375452836
Recently I solved this problem using a different observation from the official editorial, so I thought I'd share it.
Let the original array be:
After gravity shifts to the right, the final configuration is equivalent to sorting the heights in non-decreasing order. Let the sorted array be:
Instead of tracking every cube individually, we can compute the answer using weighted positional sums.
Define:
and
Then the total movement of all cubes is simply:
Therefore, once the array is fixed, the answer can be computed in O(n log n) by sorting the heights and evaluating the two sums.
We are allowed to decrease exactly one height by 1.
Suppose we decrease the height of column i. Instead of recomputing the entire answer for every position, we only need to know how much extra movement this removal creates.
Claim: The increase in total movement is exactly equal to the number of elements on its left having value greater than or equal to a_i.
Define:
Then the optimal move is to choose the index having maximum value of greater_i.
Consider the block removed from column i. Let
At height h, every column whose height is at least h contributes exactly one block.
Suppose there are several such blocks.
After gravity shifts to the right, all these blocks become consecutive in the final sorted configuration. Now remove the block belonging to column i at height h.Every participating block that originally appeared before column i now moves one position further to the right in the final arrangement.Each such block gains exactly 1 additional unit of movement.
Therefore the total increase in movement equals the number of participating blocks before column i. A column participates at height h if and only if its height is at least h.
Hence the gain becomes
which is exactly greater_i.
A naive approach checks every previous element for every position.
Complexity:
which is too slow.
We instead use:
• Coordinate Compression
• Fenwick Tree (Binary Indexed Tree)
For every position:
Compress the value to its rank.
Query how many strictly smaller values have already appeared.
Let this value be
Since there are i−1 previous elements,
Fenwick Tree supports update and query operations in
Therefore all greater_i values are computed in
Coordinate Compression:
Fenwick Tree Processing:
Sorting:
Overall Complexity:
Memory Complexity:
This approach views the problem through positional sums and uses a Fenwick Tree to determine the optimal column to decrease.
I found this observation quite intuitive because the entire effect of the operation can be reduced to counting how many previous heights are greater than or equal to the current height.
Any suggestions, alternative proofs, or improvements are welcome.
