SHFTGAME - Editorial

PROBLEM LINK:

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

Author: biggestotaku
Tester: kingmessi
Editorialist: iceknight1093

DIFFICULTY:

Medium

PREREQUISITES:

Binary search, prefix sums

PROBLEM:

There are two binary strings A and B.
Consider the following process:

  • You start with a score of 0.
  • Every turn, add A_1 \cdot B_1 to your score.
  • Then, if A_1 = B_1, left-rotate A by one step.
    Otherwise, left-rotate B by one step.

Answer Q queries: each query gives you an integer K and asks for your score after K iterations of the process.

EXPLANATION:

Note that because we work with binary strings, A_1 \cdot B_1 = 1 if and only if A_1 = B_1 = 1.
Any other combination of values will result in a product of 0.

Also, rather than left-rotation, it’s perhaps more helpful to visualize the process in terms of moving two pointers instead.
That is, suppose we have two pointers p and q. p points to an index in A, and q points to an index in B.
Initially, p = q = 1.
Then, left-rotating A is equivalent to incrementing p by 1, while left-rotating B is equivalent to incrementing q by 1.
When either p or q becomes N+1 via an increment, we set them back to 0.

With this model, the score at pointer state (p, q) is exactly A_p \cdot B_q.


Let’s now analyze how the score actually evolves from any given state.

Suppose A_p = 0 and B_q = 0.
The product is 0, and because A_p = B_q, p will increment by 1.
Note that if A_{p+1} = 0, the exact same logic holds again - the score doesn’t change and the pointer corresponding to A will just increment by 1.
So, this will keep happening till we finally reach a 1 in A.

Next, look at A_p = 1 and B_q = 0.
Again the product is 0, but now A_p \ne B_q so q will increment instead.
Once again, this will continue to happen till we reach a 1 in B.

Next, consider A_p = B_q = 1.
Here, the score increases by 1, and p will increment.
So, the score will repeatedly increase by 1 till we reach a 0 in A.

Finally, consider A_p = 0 and B_q = 1.
Here, the score doesn’t change, and q will repeatedly increase till we reach a 0 in B.


From the above analysis, we see that the process follows a somewhat regular pattern.
In particular, the values of (A_p, B_q) will cycle through (0, 0) \to (1, 0) \to (1, 1) \to (0, 1) \to (0, 0) in order.
Further, each pointer will always move through an entire ‘block’ of equal elements; and movements will always alternate between pointers: that is, first p will move through a block of equal elements, then q will move through a block of equal elements, then p again, then q again, and so on.

This motivates us to work with blocks of equal elements in both strings.

For convenience, let’s work with a “nice enough” starting state.
Specifically, let’s first perform a few moves manually and record their results, and stop exactly when we first move to a (0, 0) state from a (0, 1) state.
Note that this will certainly happen within N moves, so direct simulation (till here) is not too slow.

Suppose the state we currently are at is (p_s, q_s).
Then, we have the nice property that A_{p_s} = B_{q_s} = 0 while A_{p_s - 1} = B_{q_s - 1} = 1 (of course, considering indices cyclically if needed.)
That is, we can pretend that both strings “start with 0” and “end with 1”.
(You may note that there is an edge case, wherein A and/or B consists entirely of a single character - we’ll return to those cases at the end.)

Now that we have this “nice” starting state, let’s compute the lengths of the blocks of equal elements in both strings, starting from indices p_s and q_s, respectively.
Suppose these lengths are (x_1, x_2, \ldots, x_a) for A and (y_1, y_2, \ldots, y_b) for B.
Here, because blocks must alternate between 0 and 1, and we start with 0, the odd-indexed values will represent lengths of 0-blocks, while the even-indexed values represent lengths of 1-blocks.

Importantly, observe that a and b will be even because our strings end with 1 because we brought them into a nice starting state first.


With these block lengths in hand, let’s analyze how the pointers move and the score evolves.
From what we saw previously, since we start at (0, 0),

  • First, we’ll step through x_1 indices of A without the score increasing.
  • Second, we’ll step through y_1 indices of B without the score increasing.
  • Third, we’ll step through x_2 indices of A, with the score increasing for each move.
  • Finally, we’ll step through y_2 indices of B without the score increasing.
  • This pattern will then repeat with the indices (x_3, y_3, x_4, y_4) and so on.
    Observe that due to the cyclic nature of things, indices in x can be considered modulo a and indices in y can be considered modulo b (as in, once we reach the ‘last’ pair of blocks in A, we’ll just cycle around to the first ones irrespective of what’s happening in B.)

With this observation, we’re now ready to find the answer for some given parameter K.

First, recall that we performed a few moves manually to reach our nice starting state.
Let the number of manual moves be M.
If K \le M, we already know what the answer is; otherwise we can add the result of the first M moves to our answer and solve for (K - M) using the block decomposition.

Now, we work with the blocks.
First, let’s compute the number of “full” blocks that have been traversed.
One way to do this is to use binary search and some math/prefix sums.
Specifically, suppose we want to check if F full blocks have been traversed without crossing K moves.
Then, we must have used \text{ceil}(\frac{F}{2}) blocks from A and \text{floor}(\frac{F}{2}) blocks from B.
The sum of lengths of these blocks can be computed using prefix sums on the (x_i) and (y_i) sequences; since each sequence will be fully traversed several times and then some prefix will remain.
Since this is doable quickly, we can binary search to find the largest valid F.

Now once F is known, we can also compute the total addition to the score, with more prefix sums (this time only on the even-index elements of (x_i)).
Finally, there will be one block that’s partially traversed; if this block is an even-index x_i it will contribute a bit to the answer so add that bit too.

This allows us to solve a single query in \mathcal{O}(\log N) time after \mathcal{O}(N) precomputation, which is fast enough.


Finally, we return to the ‘special cases’ we left aside at the start, where one or both strings consist of only a single type of character.

If either A or B consist of only zeros, the answer is clearly always 0.

If A contains only ones, then observe that a small number of initial moves (no more than N) will all have a product of 0 while B’s pointer moves, and then as soon at the product becomes 1 for the first time it will remain 1 forever because only A’s pointer will move.
This case is hence easily solved by knowing the first time the product becomes 1.

If B contains only ones, we have a similar situation but in the opposite: the product will be 1 for some prefix of moves (at most N), and once it becomes 0 it will remain 0 forever.
Again this is easy to solve.

TIME COMPLEXITY:

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

CODE:

Author's code (C++)
#include<bits/stdc++.h>
#define int long long
#define nl "\n"
using namespace std;

void solve() {
    int n, q;
    cin >> n >> q;
    string A, B;
    cin >> A >> B;
    int cA[2] = {}, cB[2] = {};
    for (int i = 0; i < n; i++) {
        cA[A[i] - '0']++;
        cB[B[i] - '0']++;
    }
    if (cA[1] == 0 || cB[1] == 0) { // Any all Zeros
        while (q--) {
            int x;
            cin >> x;
            cout << 0;
            if (q) cout << " ";
            else cout << nl;
        }
        return;
    }
    if (cA[0] == 0 && cB[0] == 0) { // All ones
        while (q--) {
            int x;
            cin >> x;
            cout << x;
            if (q) cout << " ";
            else cout << nl;
        }
        return;
    }
    if (cA[0] == 0) { // A all ones
        int id = 0;
        while (B[id] != '1') id++;
        while (q--) {
            int x;
            cin >> x;
            cout << max(0ll, x - id);
            if (q) cout << " ";
            else cout << nl;
        }
        return;
    }
    if (cB[0] == 0) { // B all ones
        int id = 0;
        while (A[id] != '0') id++;
        while (q--) {
            int x;
            cin >> x;
            cout << min(x, id);
            if (q) cout << " ";
            else cout << nl;
        }
        return;
    }
    int score = 0;
    vector<int> inivals(2 * n + 2);
    int ptr = 1, pA = 0, pB = 0;
    while ((A[pA] != '0' && A[(pA + n - 1) % n] != '1') || (B[pB] != '0' && B[(pB + n - 1) % n] != '1')) {
        score += (A[pA] - '0') * (B[pB] - '0');
        if (A[pA] == B[pB]) pA = (pA + 1 % n);
        else pB = (pB + 1) % n;
        inivals[ptr++] = score;
    }
    inivals.resize(ptr);
    rotate(A.begin(), A.begin() + pA, A.end());
    rotate(B.begin(), B.begin() + pB, B.end());
    vector<int> vA, vB;
    for (int i = 0, lA = 0, lB = 0; i < n; i++) {
        lA++, lB++;
        if (i == n - 1 || A[i] != A[i + 1]) {
            vA.push_back(lA);
            lA = 0;
        }
        if (i == n - 1 || B[i] != B[i + 1]) {
            vB.push_back(lB);
            lB = 0;
        }
    }
    vector<int> a_ps_odd(vA.size() / 2);
    for (int i = 1, sm = 0; i < vA.size(); i += 2) {
        sm += vA[i];
        a_ps_odd[i / 2] = sm;
    }
    vector<int> a_ps(vA.size()), b_ps(vB.size());
    for (int i = 0, sm = 0; i < vA.size(); i++) {
        sm += vA[i];
        a_ps[i] = sm;
    }
    for (int i = 0, sm = 0; i < vB.size(); i++) {
        sm += vB[i];
        b_ps[i] = sm;
    }
    function<int(int)> steps = [&](int blcks) {
        int blckA = (blcks + 1) / 2, blckB = blcks / 2;
        int res = n * (blckA / vA.size() + blckB / vB.size());
        blckA %= vA.size();
        blckB %= vB.size();
        if (blckA) res += a_ps[blckA - 1];
        if (blckB) res += b_ps[blckB - 1];
        return res;
    };
    function<int(int)> get_ones = [&](int blcks) {
        int res = a_ps_odd.back() * (blcks / vA.size());
        blcks %= vA.size();
        if (blcks > 1) res += a_ps_odd[blcks / 2 - 1];
        return res;
    };
    while (q--) {
        int x;
        cin >> x;
        if (x < ptr) {
            cout << inivals[x];
            if (q) cout << " ";
            else cout << nl;
            continue;
        }
        int res = inivals[ptr - 1];
        x -= ptr - 1;
        int l = 0, r = x;
        while (l < r) {
            int m = (l + r + 1) / 2;
            if (steps(m) <= x) l = m;
            else r = m - 1;
        }
        int pA = (l + 1) / 2 % vA.size();
        int pB = l / 2 % vB.size();
        x -= steps(l);
        res += (pA & 1) * (pB & 1) * x;
        res += get_ones((l + 1) / 2);
        cout << res;
        if (q) cout << " ";
        else cout << nl;
    }
}

int32_t main() {
    ios_base::sync_with_stdio(0);
    cin.tie(0);
    int t = 1;
    cin >> t;
    while (t--) {
        solve();
    }
}