DIFREP - Editorial

PROBLEM LINK:

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

Author: iceknight1093
Tester: sushil2006
Editorialist: iceknight1093

DIFFICULTY:

Easy-Medium

PREREQUISITES:

Binary search

PROBLEM:

You have an array A with elements between 1 and N.
The following operation can be performed:

  • Choose L, R, X such that 1 \leq L \leq R \leq N and 1 \leq X \leq N, such that X doesn’t appear in the subarray [A_L, A_{L+1}, \ldots, A_R].
    Then, set all the elements in this subarray to X.

How many of these moves result in A being sorted?

EXPLANATION:

Suppose the subarray [L, R] is modified.
Then, note that for the array to be sorted, the prefix [1, L-1] and the suffix [R+1, N] must definitely be sorted themselves.

So, let l be the leftmost index such that A_l \gt A_{l+1}, and r be the rightmost index such that A_r \gt A_{r+1}.
These are the left/rightmost indices that break sorting, and so certainly L \leq l+1 and R \geq r must hold.
In particular, no matter what subarray is chosen, the range [l+1, r] must always be included in it.


Let’s fix the value of X, and try to figure out which subarrays [L, R] don’t contain X, but can be replaced with X to make A sorted.
First, as noted earlier, [l+1, r] must definitely be contained in the chosen subarray - so if X appears here, no valid [L, R] exists.

So, suppose X doesn’t appear in [l+1, r].
Recall that we know [1, l] and [r+1, N] are already sorted.
If the middle part is replaced by X entirely, then any part of the prefix that’s \gt X (which will be some subarray ending at l) must also be replaced by X to make A sorted. Similarly, any part of the suffix that’s \lt X, which will be some subarray starting at r+1, must be replaced too.

Let p \leq l+1 denote the first index containing an element \gt X, and q\geq r be the last index containing an element \lt X.
p and q can be found using binary search since the prefix/suffix under consideration are sorted - just make sure to treat A_{l+1} and A_r as infinity.
Then, the above observation tells us that we must have L \leq p and R \geq q.

Now, it can be seen that:

  • If A_{p-1} = X, the only choice for L is p itself: anything smaller would include an occurrence of X, which isn’t allowed.
  • If A_{p-1} \lt X instead, then L can be anything at all from 1 to L, and this choice is independent of R.
  • Similarly, depending on whether A_{q+1} = X or not, R is either fixed to be exactly q, or can be anything from q to N.

This means we can start with 1, and then multiply L and (N-q+1) to the count if the appropriate conditions are satisfied.
Either way, once p and q are known, the number of subarrays can be found in constant time.
Iterating over all X and adding up the answers gives us a solution in \mathcal{O}(N\log N).


The above analysis breaks down when A is already sorted - the algorithm works with the sorted prefix/suffix being disjoint, but in this case they’ll both equal the entire array.
However, A being sorted means it’s fairly easy to reason about sortedness anyway:

  • If X does appear in A, it’ll be in some range, say [p, q].
    Then, the only valid subarrays are all those that end at p-1 or start at q+1.
  • If X doesn’t appear in A, any valid subarray [L, R] must have A_{L-1} \lt X and A_{R+1} \gt X, it’s not too hard to count these if the breakpoint between numbers \lt X and \gt X in A is found.

TIME COMPLEXITY:

\mathcal{O}(N\log N) per testcase.

CODE:

Author's code (C++)
// #include <bits/allocator.h>
// #pragma GCC optimize("O3,unroll-loops")
// #pragma GCC target("avx2,bmi,bmi2,lzcnt,popcnt")
#include "bits/stdc++.h"
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<int> a(n);
        for (int &x : a) cin >> x;

        int L = 1, R = n-2;
        while (L < n and a[L] >= a[L-1]) ++L;
        while (R >= 0 and a[R] <= a[R+1]) --R;

        vector<int> mark(n+1), freq(n+1);
        for (int i = 0; i < n; ++i) {
            freq[a[i]] += 1;
            if (L <= i and i <= R) mark[a[i]] = 1;
        }
        
        ll ans = 0;
        for (int x = 1; x <= n; ++x) {
            if (mark[x]) continue;

            int lt = lower_bound(begin(a), begin(a) + L, x) - begin(a);
            int rt = lower_bound(begin(a) + R + 1, end(a), x) - begin(a);

            if (L == n) { // Already sorted
                if (freq[x]) {
                    ans += n - freq[x];
                }
                else {
                    ans += n - lt + 1ll * lt * (n + 1 - lt);
                }
                continue;
            }

            ll ways = 1;
            if (a[lt] != x) ways *= lt + 1;
            if (rt == n or a[rt] != x) ways *= n - rt + 1;
            ans += ways;
        }
        cout << ans << '\n';
    }
}
Tester's code (C++)
#include <bits/stdc++.h>
#include <ext/pb_ds/assoc_container.hpp>
#include <ext/pb_ds/tree_policy.hpp>

using namespace std;
using namespace __gnu_pbds;

template<typename T> using Tree = tree<T, null_type, less<T>, rb_tree_tag, tree_order_statistics_node_update>;
typedef long long int ll;
typedef long double ld;
typedef pair<int,int> pii;
typedef pair<ll,ll> pll;

#define fastio ios_base::sync_with_stdio(false); cin.tie(NULL)
#define pb push_back
#define endl '\n'
#define sz(a) (int)a.size()
#define setbits(x) __builtin_popcountll(x)
#define ff first
#define ss second
#define conts continue
#define ceil2(x,y) ((x+y-1)/(y))
#define all(a) a.begin(), a.end()
#define rall(a) a.rbegin(), a.rend()
#define yes cout << "YES" << endl
#define no cout << "NO" << endl

#define rep(i,n) for(int i = 0; i < n; ++i)
#define rep1(i,n) for(int i = 1; i <= n; ++i)
#define rev(i,s,e) for(int i = s; i >= e; --i)
#define trav(i,a) for(auto &i : a)

template<typename T>
void amin(T &a, T b) {
    a = min(a,b);
}

template<typename T>
void amax(T &a, T b) {
    a = max(a,b);
}

#ifdef LOCAL
#include "debug.h"
#else
#define debug(...) 42
#endif

/*



*/

const int MOD = 1e9 + 7;
const int N = 1e5 + 5;
const int inf1 = int(1e9) + 5;
const ll inf2 = ll(1e18) + 5;

void solve(int test_case){
    ll n; cin >> n;
    vector<ll> a(n+5);
    rep1(i,n) cin >> a[i];

    a[0] = 1, a[n+1] = n;
    ll L = n, R = 1;

    rep1(i,n){
        if(a[i] < a[i-1]){
            L = i;
            break;
        }
    }

    rev(i,n,1){
        if(a[i] > a[i+1]){
            R = i;
            break;
        }
    }

    vector<ll> pos[n+5];
    rep1(x,n) pos[x].pb(0);
    rep1(i,n) pos[a[i]].pb(i);
    rep1(x,n) pos[x].pb(n+1);

    ll ans = 0;

    rep1(x,n){
        auto &p = pos[x];
        rep(i,sz(p)-1){
            ll lx = p[i]+1, rx = p[i+1]-1;
            if(lx > rx) conts;

            // find max l
            // a[l-1] <= x

            ll mxl = -1;
            {
                ll lo = lx, hi = min(L,rx);
                while(lo <= hi){
                    ll mid = (lo+hi)>>1;
                    if(a[mid-1] <= x){
                        mxl = mid;
                        lo = mid+1;
                    }
                    else{
                        hi = mid-1;
                    }
                }
            }

            // find min r
            // a[r+1] >= x

            ll mnr = -1;
            {
                ll lo = max(R,lx), hi = rx;
                while(lo <= hi){
                    ll mid = (lo+hi)>>1;
                    if(a[mid+1] >= x){
                        mnr = mid;
                        hi = mid-1;
                    }
                    else{
                        lo = mid+1;
                    }
                }
            }

            if(mxl == -1 or mnr == -1) conts;

            ll ways = (mxl-lx+1)*(rx-mnr+1);
            if(mxl > mnr){
                ll c = mxl-mnr;
                ways -= c*(c+1)/2;
            }
            
            ans += ways;
        }
    }

    cout << ans << endl;
}

int main()
{
    fastio;

    int t = 1;
    cin >> t;

    rep1(i, t) {
        solve(i);
    }

    return 0;
}

i was fixing L or R, fixing X never occurred in my mind, fixing X made it easy.
Thanks!