TRANSFER - Editorial

editorial
, ,
1

PROBLEM LINK:

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

Author: raysh07
Tester: iceknight1093
Editorialist: iceknight1093

DIFFICULTY:

Easy-medium

PREREQUISITES:

None

PROBLEM:

You’re given an array A.
You can choose an index i with A_i \gt A_{i+1} and subtract 1 from A_i, add 1 to A_{i+1}.
Find the maximum number of operations that can be done.

EXPLANATION:

Performing an operation on an index i both reduces A_i and increases A_{i+1}, which in turn makes it easier for those indices to be involved in other operations (for indices i-1, i+1 respectively.)
However, it’s not really clear which order we must perform operations in.

To understand this, we view the operation on a slightly transformed array: the prefix sums of A.
Let P_i = A_1 + \ldots + A_i denote the i-th prefix sum of A.

Operating on index i in A is then equivalent to subtracting 1 from P_i, which is a somewhat simpler operation to reason about.
However, we aren’t free to subtract however we like: an operation can only be performed when A_i \gt A_{i+1}, which translates to P_i - P_{i-1} \gt P_{i+1} - P_i, or equivalently 2P_i \gt P_{i-1} + P_{i+1}
(this assumes P_0 = 0.)

So, we have the array P, and we can reduce its elements one at a time (under certain conditions), with our goal being to maximize the number of reductions.
The “end goal” is when we reach a state where 2P_i \le P_{i-1} + P_{i+1} for all i.

Looking at things from this perspective brings out a rather interesting property: the final array is unique!
More precisely, it doesn’t matter what order we operate in - we’ll always reach the same end state!

Proof

Let C be an array such that:

  1. C_0 = 0
  2. C_N = P_N
  3. C_i \le P_i for all i.
  4. 2C_i \le C_{i-1} + C_{i+1} for each 1 \le i \lt N.

C is one potential final array for us, since P_0 and P_N can’t be changed; while everything else can be decreased by 1.

We’ll prove that we can never go “below” C no matter what we do, i.e. it’s never possible to make index i have value smaller than C_i.

Suppose we’ve performed several operations, and the current array is Y.
Also suppose that we’re able to perform a move on some index i that will make Y_i \lt C_i after it; and this is the first time that such a situation is happening.
That means we must currently have Y_i = C_i, and also Y_j \ge C_j for all j.

However,

  • Being able to perform a move at Y_i implies that 2Y_i \gt Y_{i-1} + Y_{i+1}.
  • This means 2C_i \gt Y_{i-1} + Y_{i+1}.
  • This in turn means 2C_i \gt C_{i-1} + C_{i+1}.
  • But this contradicts the fact that 2C_i \le C_{i-1} + C_{i+1} for all i.

So, it’s never possible for us to perform a move that puts us below C at any point at all.

Now, C was just an arbitrary array that satisfied the given properties; and so this result holds for the “largest” such C as well.
Interestingly, because the endpoints are fixed (C_0 = 0 and C_N = P_N), there is in fact a unique “largest” such array!
Here, “largest” means that there’s an array M satisfying the properties, and also C_i \le M_i for every index for any other valid array C.

This follows from the fact that the array C is really just a convex array, so we are in some sense taking a convex hull.

The fact that the final state is unique is very helpful, because it means that if we’re able to simply simulate the process quickly enough, we’ll end up at the requisite end state; after which it’s easy to compute the answer (if the final array we reach is B, then the answer is the sum of prefix sums of A minus the sum of prefix sums of B, since each move reduces a prefix sum by exactly 1.)

One way to do this is to simply process the array from left to right.

Suppose we’ve processed the first i elements, and after performing all possible operations they’ve stabilized into B_1 \le\ldots\le B_i.
Let’s see what happens when we then append a new element A_{i+1} to this.

If B_i \le A_{i+1}, nothing really happens - no new operation is possible, so we can just set B_{i+1} = A_{i+1} and continue on.

But, if B_i \gt A_{i+1}, some operations are possible.
If we try to perform them, it’ll end up something like this:

  • Set B_{i+1} = A_{i+1} first.
  • Try to transfer from B_i to B_{i+1} as much as possible.
    This will cause the values to become equal (or differ by 1, if their sum is odd.)
  • After this, the value of B_i has reduced, so it might be possible to perform some operations using (B_{i-1}, B_i).
  • After those are done, it might be possible to perform some more operations with (B_{i-2}, B_{i-1}) and also (B_i, B_{i+1}) since B_i might be larger now.
    \vdots

So it might be that we’ll have to operate on the same index multiple times; there’s no simple direct simulation.

However, we don’t need to do that!
We can instead directly find out the final state by seeing what it’ll look like.

Observe that no matter what, only some suffix of B (ending at i+1) will change in value.
Say, the elements B_j, B_{j+1}, \ldots, B_{i+1} end up changing.
What will their values be?

Well, if we let S = B_j + \ldots + B_{i+1} be their total sum, and C = i-j+2 be their count, then due to the uniqueness of the solution (and ending up at the “greatest possible” non-decreasing sequence), these elements will be as equal as possible; meaning they’ll differ by at most 1.
Hence,

  • If C divides S, all these C elements will become \frac{S}{C}.
  • Otherwise, let q = \left\lfloor \frac{S}{C} \right\rfloor and r = S \bmod C.
    There will be C-r copies of q and r copies of q+1, in sorted order naturally.

So, we only really need to find this index j.

It is now in fact possible to solve the problem already, by throwing some machinery at it.
The index j can be found using binary search (along with a range sum data structure), after which you can use a range-set data structure to perform appropriate updates.
This is \mathcal{O}(N \log^2 N) or \mathcal{O}(N\log N) depending on implementation.
The \mathcal{O}(N \log^2 N) in particular may not pass depending on constant factor.

However, there is a far better approach - it’s possible to solve the problem in linear time, without any heavy data structures at all!

To do this, we’ll represent the array as blocks of equal elements. The blocks will be non-decreasing in value from left to right.
Consider what happens when we append a new element (which is its own singleton block.)

Looking at the last two blocks, if they’re not in sorted order we can merge them together.
As seen earlier, merging them will create (at most) two blocks; which can be determined from the sum and count of elements.

We can then check if this results in sorted order - and if not, bring the third-last block into the fold as well. This can be done repeatedly till we end up at sorted order.

While this looks like it might take quadratic time, if implemented properly it’s actually (amortized) linear.
This is because eventually, each element we append will create at most two new blocks; and each block can be deleted only once so we create at most 2N blocks and delete each of them at most once, all in constant time per operation.

This allows us to build the final array in terms of blocks, after which computing the actual number of operations is trivial (as mentioned earlier, the difference between the sum of prefix sums of the initial and final arrays.)

TIME COMPLEXITY:

\mathcal{O}(N) per testcase.

CODE:

Editorialist's code (C++) 
// #pragma GCC optimize("O3,unroll-loops")
// #pragma GCC target("avx2,bmi,bmi2,lzcnt,popcnt")
#include "bits/stdc++.h"
#include <cassert>
using namespace std;
using ll = long long int;
mt19937_64 RNG(chrono::high_resolution_clock::now().time_since_epoch().count());

int main()
{
    ios::sync_with_stdio(false); cin.tie(0);

    int t; cin >> t;
    while (t--) {
        int n; cin >> n;
        vector a(n, 0);
        for (int &x : a) cin >> x;

        vector<array<ll, 2>> blocks = {{0, 0}};
        auto push = [&] (const auto &self, ll x, ll k) -> void {
            ll sm = x*k, ct = k;
            while (true) {
                ll val = sm/ct;
                if (val < blocks.back()[0]) {
                    auto [y, c] = blocks.back();
                    blocks.pop_back();
                    sm += y*c;
                    ct += c;
                }
                else break;
            }
            
            int r = sm % ct;
            blocks.push_back({sm/ct, ct - r});
            if (r > 0) blocks.push_back({sm/ct + 1, r});
        };

        for (int x : a) push(push, x, 1);

        vector<ll> b;
        for (auto [x, y] : blocks) {
            for (int i = 0; i < y; ++i) b.push_back(x);
        }

        ll ans = 0, p1 = 0, p2 = 0;
        for (int i = 0; i < n; ++i) {
            p1 += a[i];
            p2 += b[i];
            ans += p1 - p2;
        }
        cout << ans << '\n';
    }
}