EQAR - Editorial

PROBLEM LINK:

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

Author: helloLad
Tester: wasd2401
Editorialist: iceknight1093

DIFFICULTY:

TBD

PREREQUISITES:

Greedy

PROBLEM:

You’re given an array A of length N.

Define F(A) as follows:
In one move, you can pick indices 1 \leq i_1 \lt i_2 \lt\ldots\lt i_k \leq N such that i_j+1 \neq i_{j+1}, i.e, a set of non-adjacent indices.
Then, add 1 to each A_{i_j}, and subtract 1 from some element between each adjacent pair of indices; or vice versa.

F(A) denotes the the minimum number of operations needed to make all the elements of A equal.

Given A, find the sum of F(A[1\ldots i]) across all prefixes of A.

EXPLANATION:

First, let’s solve the problem of computing F(A) for a single array.

Note that operation can be written more simply as, “choose any odd-length subsequence of A, then add +1 and -1 to alternate elements of this subsequence”.

Suppose we want to make all the elements of A equal x.
Then, we can make a couple of observations about the structure of an optimal solution:

  • If A_i \lt x, there will exist an optimal solution where every move that changes A_i, only increases it.
    Similarly, if A_i \gt x, there exists an optimal solution that only ever decreases A_i.
    And of course, if A_i = x we can ensure that A_i is never changed.
    This holds simultaneously for all indices.
  • Let B_i = X - A_i denote the amount that A_i needs to change by.
    The above observation tells us that every move involves choosing some odd-length subsequence of B such that adjacent elements have opposite signs, and adding/subtracting 1 appropriately.
    A ‘natural’ greedy solution is then as follows:
    • Choose the leftmost non-zero element of B; wlog let it be positive, and say this is index i_1.
    • Choose the smallest i_2\gt i_1 such that B_{i_2} \lt 0.
    • Choose the smallest i_3\gt i_2 such that B_{i_3} \gt 0.
      \vdots
    • Choose the smallest i_k\gt i_{k-1} such that B_{i_k} \gt 0.
    • Operate on this subsequence, then repeat till B becomes filled with zeros.

It turns out that this greedy is indeed correct!
A proof of this (and the first claim) can be found at the bottom.


The greedy idea, while correct, is still much too slow to solve this problem, so we need to speed it up.
To do that, we can think of ‘extending’ existing operations.

That is, suppose we’re at index i, such that we’ve already decided the moves for indices \lt i: more specifically, we’ve decided the prefix of each move till i-1.
That means each move so far is one of four types: either even or odd length, and ending with either +1 or -1.
Let their counts be c_{00}, c_{01}, c_{10}, c_{11} respectively, where the first digit denote the parity of length, and the second denotes whether it ends with +1 or -1 (so c_{10} is odd-length moves ending with +1, for instance).

Suppose B_i \gt 0. Then,

  • We can take an even-length move ending with -1 and extend it to form an odd-length move.
    • This can be done upto \min(B_i, c_{01}) times.
      Note that all these moves will turn into odd-length moves ending with +1 now.
      Make sure to update B_i appropriately too.
  • We can take an odd-length move ending with -1 and extend it to form an even-length move.
    • Again, this can be done upto \min(B_i, c_{11}) times.
      Note that this doesn’t create ‘legitimate’ moves, since they’ll be even-length after it; but that’s ok: they can be legitimized in the future (and if they can’t, we’ll take care of that too at the end).
  • If B_i \gt 0 still, our only option is to start new moves at this index; so we add B_i to c_{10}.

If B_i \lt 0, the cases are similar.

Now, let’s look at the values of c_{ij} after all the indices have been processed.

  • There will certainly be c_{10} + c_{11} moves used; being odd-length.
  • We also have c_{00} + c_{01} moves that are even-length.
    Note that each of them can be turned into odd-length moves by just removing their last element and making that a single move; and we can’t do any better, as per our greedy choice.
    So, these contribute 2\cdot (c_{00} + c_{01}) moves in total.

So, the total number of moves required is c_{10} + c_{11} + 2\cdot (c_{00} + c_{01}).

Notice that in the process of computing these values, we in fact computed them for every prefix of A.
So, for a fixed x, we can compute the minimum number of moves to turn each prefix of A into x in linear time.

So, do this for each x from 1 to 5000 (or from \min(A) to \max(A), for slightly better bounds), and take the best answer for each prefix; then add them up to finish.


Proofs

Claim: Fix x, and compute the array B_i = x - A_i.
Then, there exists an optimal solution that only brings each B_i closer to 0.

Proof: Let B_i \geq 0, and suppose we perform a move that increases it.
Let the indices of this move be x_1, x_2, \ldots, x_r, \ldots, where x_r = i.

Since B_i must definitely become 0 in the end, there will exist another move that decreases B_i.
Let this be y_1, y_2, \ldots, y_s, \ldots with y_s = i.

Then, we can do the following:

  • Look at x_{r-1} and y_{s-1}. Suppose they both exist.
    • If x_{r-1} = y_{s-1}, delete x_{r-1}, x_r from the first move and y_{s-1}, y_s from the second.
      This won’t change the end result nor the moves at other indices, and we now have one less move that increases index i.
    • If x_{r-1} \lt y_{s-1}, delete x_r from the first move and y_s from the second, and move y_{s-1} to the first move. This ensures that both moves remain valid (i.e have odd length), and doesn’t change their overall effect.
    • If x_{r-1} \gt y_{s-1}, do the same but instead move x_{r-1} to the second move.
  • If x_{r-1} doesn’t exist but y_{s-1} does, delete x_r and y_s from the moves, and switch y_{s-1} to the first move (just as was done above).
    If $x_{r-1} exists and y_{s-1} doesn’t, the situation is similar.
  • If both x_{r-1} and y_{s-1} don’t exist, look at x_{r+1} and y_{s+1} instead; and perform similar casework.
  • If none of these four indices exist, we know that both moves operate on the singleton index i; so just don’t perform both moves.

By repeatedly doing this, we can ensure that elements with B_i \geq 0 are only reduced, and those with B_i \leq 0 are only increased till they reach 0.
This applies to every index simultaneously, since when fixing the moves affecting one index, we didn’t switch the type of move being applied at any other index.


Claim: There exists an optimal solution that repeatedly greedily chooses moves from left to right.

Proof: Consider any optimal sequence of moves. Their order doesn’t matter, so wlog we can assume the first move changes the leftmost non-zero element.

Let i_1 be the leftmost non-zero element, wlog say B_{i_1} \gt 0.
if there are no negative elements in B, clearly the only choice is to decrease B_{i_1}, so this is optimal.

Otherwise, let the sequence of indices chosen be i_1, x_2, x_3, \ldots
Let i_2 denote the first negative element after i_1.
Suppose x_2 \neq i_2, which also means x_2 \gt i_2.
Then, there must exist some other move that changes i_2. Let the index immediately following i_2 in this move be y.
Then,

  • If y \gt x_2, simply move i_2 to the first move and x_2 to the second.
  • If y \lt x_2, we have i_1 \lt i_2 \lt y \lt x_2, so move both i_2 and y to the first move.

There’s still one issue to deal with: y and/or x_2 might not exist.

  • If y doesn’t exist and x_2 does, move i_2 to the first move and operate on x_2 alone.
  • If y exists and x_2 doesn’t, move both y and i_2 to the first move.
  • If both don’t exist, then there are two possibilities:
    • There doesn’t exist a positive element after i_2. Then, our only choice is operating on them separately anyway.
    • If there does exist a positive element after i_2, we have more cases:
      1. If some move starts at a positive element after i_2, both i_1 and i_2 can be included in it.
      2. If some move starts at a negative position \gt i_2, replace its starting element with i_2 instead; and now y exists so go back to those cases.
      3. If both the above fail, there must be a move that starts at a positive position between i_1 and i_2, so replace its start with i_1 and now x_2 exists.

This way, it’s always possible to include both i_1 and i_2 in the same move.
Repeat this argument for the first positive element after i_2, and so on.
Then, make this move, and repeat the argument for the remaining array.

TIME COMPLEXITY:

\mathcal{O}(N^2) per testcase.

CODE:

Author's code (C++)
#include <bits/stdc++.h>
using namespace std;
#define IOS ios_base::sync_with_stdio(0); cin.tie(0); cout.tie(0);
#define int long long
int dp[5001][5001][5];
// vector<int> dp[5001][5001];

signed main(){
    IOS
    #ifndef ONLINE_JUDGE
    freopen("input.txt", "r", stdin);
    freopen("output.txt", "w", stdout);
    #endif
    int t;
    cin>>t;
    while(t--){
        int n;
        cin>>n;
        vector<int> a(n+1);
        for(int i=1;i<=n;++i){
            cin>>a[i];
        }
        for(int i=1;i<=n;++i){
            for(int j=0;j<5001;++j){
                int inc1=0,dec1=0,inc2=0,dec2=0,res=0;
                if(i>1){
                    inc1=dp[i-1][j][0],inc2=dp[i-1][j][1],
                    dec1=dp[i-1][j][2],dec2=dp[i-1][j][3],res=dp[i-1][j][4];
                }
                int val=a[i];
                if(val<j){
                    int op1=min(j-val,inc1);
                    val+=op1;
                    inc1-=op1;
                    dec2+=op1;
                    int op2=min(j-val,inc2);
                    val+=op2;
                    inc2-=op2;
                    dec1+=op2;
                    if(val<j){
                        dec2+=j-val;
                        res+=j-val;
                    }
                }else{
                    int op1=min(val-j,dec1);
                    val-=op1;
                    dec1-=op1;
                    inc2+=op1;
                    int op2=min(val-j,dec2);
                    val-=op2;
                    dec2-=op2;
                    inc1+=op2;
                    if(val>j){
                        inc2+=val-j;
                        res+=val-j;
                    }
                }
                // dp[i][j]={inc1,inc2,dec1,dec2,res};
                dp[i][j][0]=inc1,dp[i][j][1]=inc2,dp[i][j][2]=dec1,
                dp[i][j][3]=dec2,dp[i][j][4]=res;
            }
        }
        int ans=0;
        for(int i=1;i<=n;++i){
            int mn=1e18;
            for(int j=0;j<5001;++j){
                mn=min(mn,dp[i][j][0]+dp[i][j][2]+dp[i][j][4]);
            }
            ans+=mn;
        }
        cout<<ans<<'\n';
    }
    return 0;
}


Editorialist's code (Python)
for _ in range(int(input())):
    n = int(input())
    a = list(map(int, input().split()))
    ans = [10**9] * n
    for x in range(5001):
        plus1, plus2 = 0, 0 # 1 -1 1 -1 ... odd/even length
        minus1, minus2 = 0, 0 # -1 1 -1 1 ... odd/even length
        b = a[:]
        for i in range(n):
            if b[i] < x:
                take = min(x - b[i], plus2) # convert even length to odd, legitimizing the ops
                b[i] += take
                plus2 -= take
                plus1 += take

                # update answer
                ans[i] = min(ans[i], plus1 + minus1 + 2*plus2 + 2*minus2 + x - b[i])

                take = min(x - b[i], minus1) # convert odd to even if really necessary; doesn't add any new operations
                b[i] += take
                minus1 -= take
                minus2 += take

                plus1 += x - b[i] # if anything's left, no choice but to start new ops with it
            else:
                take = min(b[i] - x, minus2) # convert even length to odd, legitimizing the ops
                b[i] -= take
                minus2 -= take
                minus1 += take

                # update answer
                ans[i] = min(ans[i], plus1 + minus1 + 2*plus2 + 2*minus2 + b[i] - x)

                take = min(b[i] - x, plus1) # convert odd to even if really necessary; doesn't add any new operations
                b[i] -= take
                plus1 -= take
                plus2 += take

                minus1 += b[i] - x # if anything's left, no choice but to start new ops with it
    print(sum(ans))