STDIF - Editorial

editorial
, , , ,
1

PROBLEM LINK:

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

Author: iceknight1093
Tester: sushil2006
Editorialist: iceknight1093

DIFFICULTY:

Easy-Med

PREREQUISITES:

Binary search, hashing

PROBLEM:

An array A consisting of ones and twos is said to be stable if the following holds:

  • Consider an array B of the same length as A, initially filled with zeros.
  • You can choose an index i, and:
    • If A_i = 1, add 1 to B_i.
    • If A_i = 2, add 1 to B_{i-1}, B_i, B_{i+1}.
  • If B can be made a palindrome using these operations, A is stable.

Given an array A, count the number of its subarrays that are stable.
N \leq 2\cdot 10^5 in this version.

EXPLANATION:

Let’s recap the ad-hoc solution to this problem.
An array A will be stable if and only if at least one one of these is true:

  1. N is odd or N = 2.
  2. A_i = A_{N+1-i} for some index i.
  3. A contains three consecutive equal elements.
  4. A contains [1, 1, 2, 2] as a subarray.
  5. A starts with [1, 1] or [1, 2, 2] (or ends with their reverse).

Now though, we cannot iterate through all subarrays of A.

Let’s first take care of the odd case separately.
There are N-L+1 subarrays of length L, so the number of odd-length subarrays equals

N + (N-2) + (N-4) + \ldots + (N\bmod 2)

With this out of the way, let’s focus on the even-length subarrays.

The “standard” approach for problems where all subarrays satisfying some condition is to fix one endpoint and count valid choices for the other endpoint.
However, the hard part here is really the check for two opposing equal elements, since if we fix a left endpoint this can blink in and out of existence as the right endpoint moves.

To deal with this, we won’t fix an endpoint: instead, we’ll fix the middle of the subarray.
That is, fix an index i, and we’ll attempt to count all valid even-length subarrays centered about (i, i+1).

This can be done in a series of steps.
First, let’s find the smallest integer k_1 such that A[i-k_1 : i] contains three consecutive elements.
This is easy to do: for example store all occurrences of [1, 1, 1] or [2, 2, 2] and find the closest one with binary search, or just keep track of the nearest occurrence of both while iterating i left to right.

Similarly, let’s find the smallest integer k_2 such that A[i-k_2:i+2] contains either [1, 1, 2, 2] or [2, 2, 1, 1] as subarrays.
This can be done in the same way as the previous check.

Next, let’s find k_3 such that A_{i-k_3} = A_{i+1+k_3}, i.e. the smallest subarray centered at (i, i+1) with equal opposite elements.
This can be done in several ways, perhaps the simplest is to use binary search and hashing.
Specifically, binary search to find the largest integer k_3 such that the subarrays A[i-k_3:i] and A[i+1:i+1+k_3] have no equal opposite elements.
For a fixed k_3, this can be checked by comparing subarray hashes - you’ll need to build rolling hashes for the array as well as the reverse flipped array (i.e. 1 replaced with 2 and vice versa).

Now, let k = \min(k_1, k_2, k_3).
Note that any subarray centered at (i, i+1) with a length that’s \gt 2k will definitely satisfy one of the three conditions above; and so will surely be stable.
Thus, we can simply add the number of such subarrays to the answer.

That leaves subarrays with length \leq 2k centered about (i, i+1).
Of them, length 2 subarrays are always stable, so add 1 to the answer (if k \gt 1, that is).
For everything else, the only remaining conditions that can possibly hold, are the boundary conditions: that is, the string starting with [1, 1] or [1, 2, 2] (or ending with their reverse).
This corresponds to just counting the number of times [1, 1] and/or [1, 2,2] appear as subarrays in some range, which is again easy to do.

This allows us to process a single midpoint in \mathcal{O}(\log N) time, so we’re done.

TIME COMPLEXITY:

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

CODE:

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 = 2e5 + 5;
const int inf1 = int(1e9) + 5;
const ll inf2 = ll(1e18) + 5;

vector<int> ps;
vector<int> mods = {791139137, 433567597, 271149607, 561259969, 222708581};
ll pows[N][2];

void precalc() {
    mt19937 rng(chrono::steady_clock::now().time_since_epoch().count());
    shuffle(all(mods), rng);
    
    const int alpha = 35;
    rep(j,5){
        ps.pb(uniform_int_distribution<>(alpha+5,mods[j]-1)(rng)); 
    }

    rep(j, 2) {
        pows[0][j] = 1;
        rep1(i, N - 1) {
            pows[i][j] = pows[i - 1][j] * ps[j] % mods[j];
        }
    }
}

typedef array<ll,2> Hash;

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

    vector<Hash> fhash(n+5), rhash(n+5);
    rep1(i,n){
        ll x = a[i];
        rep(j,2){
            fhash[i][j] = (fhash[i-1][j]*ps[j]+x)%mods[j];
        }
    }

    rev(i,n,1){
        ll x = a[i];
        x ^= 3; // flip guy
        rep(j,2){
            rhash[i][j] = (rhash[i+1][j]*ps[j]+x)%mods[j];
        }
    }

    auto get_fhash = [&](ll l, ll r){
        ll len = r-l+1;
        Hash curr = fhash[r];
        Hash sub = fhash[l-1];
        rep(j,2){
            curr[j] -= sub[j]*pows[len][j];
            curr[j] = (curr[j]%mods[j]+mods[j])%mods[j];
        }
        return curr;
    };

    auto get_rhash = [&](ll l, ll r){
        ll len = r-l+1;
        Hash curr = rhash[l];
        Hash sub = rhash[r+1];
        rep(j,2){
            curr[j] -= sub[j]*pows[len][j];
            curr[j] = (curr[j]%mods[j]+mods[j])%mods[j];
        }
        return curr;
    };

    vector<bool> same_alt(n+5); // a a b b
    rep1(i,n-3){
        if(a[i] == a[i+1] and a[i+1] != a[i+2] and a[i+2] == a[i+3]){
            same_alt[i] = 1;
        }
    }

    // find closest [a a b b] to right
    vector<ll> closest_same_alt(n+5,inf2);
    rev(i,n,1){
        if(same_alt[i]){
            closest_same_alt[i] = i+3; // stores right endpoint
        }
        else{
            closest_same_alt[i] = closest_same_alt[i+1];
        }
    }

    vector<bool> is_one_triplet(n+5);
    rep1(i,n-2){
        if(a[i] == 1 and a[i+1] == 1 and a[i+2] == 1){
            is_one_triplet[i] = 1;
        }
    }

    vector<ll> closest_one_triplet(n+5,inf2);
    rev(i,n,1){
        if(is_one_triplet[i]){
            closest_one_triplet[i] = i+2;
        }
        else{
            closest_one_triplet[i] = closest_one_triplet[i+1];
        }
    }

    vector<bool> is_alt_triplet(n+5); // a b a
    rep1(i,n-2){
        if(a[i] != a[i+1] and a[i+1] != a[i+2]){
            is_alt_triplet[i] = 1;
        }
    }

    vector<ll> pref_alt_triplet(n+5);
    rep1(i,n) pref_alt_triplet[i] = pref_alt_triplet[i-1]+is_alt_triplet[i];

    auto ok = [&](ll l, ll r) -> bool{
        if(get_fhash(l,r) != get_rhash(l,r)) return 1;
        if(closest_same_alt[l] <= r) return 1;
        if(closest_one_triplet[l] <= r) return 1;
        return 0;
    };

    ll ans = 0;
    
    // +1 for all odd-size ranges
    for(int len = 1; len <= n; len += 2){
        ll cnt = n-len+1;
        ans += cnt;
    }

    // +1 for len = 2
    ans += n-2+1;

    rep1(i,n){
        // midpoint: [i,i+1]
        ll left_len = i, right_len = n-(i+1)+1;
        ll mx_extend = min(left_len, right_len);
        if(mx_extend < 2) conts;

        if(a[i] != a[i+2] or a[i+1] != a[i-1] or a[i] == a[i+1]){
            // all ranges good
            ans += mx_extend-1;
            conts;
        }

        ll lo = 1, hi = mx_extend;
        ll first_good = -1;

        while(lo <= hi){
            ll mid = (lo+hi)>>1;
            if(ok(i-mid+1,i+1+mid-1)){
                first_good = mid;
                hi = mid-1;
            }
            else{
                lo = mid+1;
            }
        }

        if(first_good != -1){
            ans += mx_extend-first_good+1;
            amin(mx_extend,first_good-1);
        }

        // if alt triplet, then bad
        ll lx = i-mx_extend+1, rx = i-1;
        if(lx <= rx){
            ll tot = rx-lx+1;
            ll bad = pref_alt_triplet[rx]-pref_alt_triplet[lx-1];
            ans += tot-bad;
        }
    }

    cout << ans << endl;
}

int main()
{
    precalc();

    fastio;

    int t = 1;
    cin >> t;

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

    return 0;
}
1 Like