REVALTCT - 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

PREREQUISITES:

None

PROBLEM:

A binary string is called almost alternating if it can be made alternating by reversing at most one of its substrings.
Given a binary string S, count the number of almost alternating substrings of S.

EXPLANATION:

First, recall the characterization of almost alternating substrings: this can be seen in the editorial to the easy version.
To recap, S is almost alternating if any one of the following conditions holds:

  • S is alternating.
  • S contains one bad index i, and either S_1 \neq S_i or S_N \neq S_i holds.
  • S contains two bad indices i and j, and S_i \neq S_j.

i is a bad index if S_i = S_{i+1}.

Let’s fix the left endpoint L of the substring that’s being considered, and try to count all valid right endpoints R.
As noted above, an almost alternating substring should contain at most two bad indices.

Let i_1, i_2, i_3 be the next three bad indices after L, with L \leq i_1 \lt i_2 \lt i_3.

Then, note that only right endpoints \leq i_3 must be considered - any further and the substring will have three bad indices.
Analyzing them,

  • For L \leq R \leq i_1, S[L, R] has no bad indices and so is alternating.
    There are i_1 - L + 1 such substrings.
  • For i_1 \lt R \leq i_2, S[L, R] has exactly one bad index, which is i_1.
    Such substrings will be almost alternating if and only if either S_L \neq S_{i_1}, or S_R \neq S_{i_1}.
    • If S_L \neq S_{i_1}, then all R in this range are valid - there are i_2 - i_1.
    • Otherwise, since characters alternate between zeros and ones from i_1 + 1 to i_2, exactly half the right endpoints there will be valid - \left\lfloor \frac{i_2 - i_1}{2} \right\rfloor to be precise.
  • For i_2 \lt R \leq i_3, S[L, R] will have exactly two bad indices.
    If S_{i_1} \neq S_{i_2} all right endpoints in this range will be valid; otherwise none of them will be.

So, once i_1, i_2, i_3 are known, the number of valid right endpoints can be computed in constant time.
There are several ways to find these three indices quickly: a simple one is to store all bad indices in a list and binary search on it, alternately you can use two pointers or even precompute the next bad index of each index.

Note that one or more of i_1, i_2, i_3 might not exist which needs to be taken care of appropriately.
A very simple way of dealing with this is to simply set any values that don’t exist to N. It can be verified that this is consistent with any formulas we use, and avoids having to deal with annoying casework.

TIME COMPLEXITY:

\mathcal{O}(N) or \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;
        string s; cin >> s;

        set<int> bad;
        for (int i = 0; i+1 < n; ++i)
            if (s[i] == s[i+1])
                bad.insert(i);
        bad.insert(n-1);
        bad.insert(n);
        bad.insert(n+1);
        ll ans = 0;
        for (int i = 0; i < n; ++i) {
            auto it = bad.lower_bound(i);
            auto it2 = bad.lower_bound(*it + 1);
            auto it3 = bad.lower_bound(*it2 + 1);

            int j = *it;
            int k = min(n-1, *it2);
            int l = min(n-1, *it3);
            
            ans += j - i + 1;            
            if (s[i] != s[j]) ans += k - j;
            else ans += (k - j) / 2;

            if (s[j] == s[k]) continue;
            ans += l - k;
        }
        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;
    string s; cin >> s;
    s = "$" + s;
    vector<ll> segs;
    rep1(i,n){
        if(i > 1 and s[i] != s[i-1]){
            segs.back()++;
        }
        else{
            segs.pb(1);
        }
    }

    ll ans = 0;

    // only 1 seg => any len is fine
    trav(x,segs){
        ans += x*(x+1)/2;
    }

    // 2 segs => at least one of them must have even len
    rep(i,sz(segs)-1){
        ll x = segs[i], y = segs[i+1];
        ans += x*y;
        ll both_odd = ceil2(x,2)*ceil2(y,2);
        ans -= both_odd;
    }

    // 3 segs => middle must have even len
    rep1(i,sz(segs)-2){
        if(segs[i]&1) conts;
        ll x = segs[i-1], y = segs[i+1];
        ans += x*y;
    }

    cout << ans << endl;
}

int main()
{
    fastio;

    int t = 1;
    cin >> t;

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

    return 0;
}