BinaryStringForces
Contest: CodeTON Round 3 (Div. 1 + Div. 2, Rated, Prizes!)
Difficulty rating: 3500
Tags: constructive algorithms, data structures, dp

You are given a binary string $$$s$$$ of length $$$n$$$. We define a maximal substring as a substring that cannot be extended while keeping all elements equal. For example, in the string $$$11000111$$$ there are three maximal substrings: $$$11$$$, $$$000$$$ and $$$111$$$.

In one operation, you can select two maximal adjacent substrings. Since they are maximal and adjacent, it's easy to see their elements must have different values. Let $$$a$$$ be the length of the sequence of ones and $$$b$$$ be the length of the sequence of zeros. Then do the following:

- If $$$a \ge b$$$, then replace $$$b$$$ selected zeros with $$$b$$$ ones.
- If $$$a < b$$$, then replace $$$a$$$ selected ones with $$$a$$$ zeros.

As an example, for $$$1110000$$$ we make it $$$0000000$$$, for $$$0011$$$ we make it $$$1111$$$. We call a string being good if it can be turned into $$$1111...1111$$$ using the aforementioned operation any number of times (possibly, zero). Find the number of good substrings among all $$$\frac{n(n+1)}{2}$$$ non-empty substrings of $$$s$$$.

Input format:
Each test consists of multiple test cases. The first line contains a single integer $$$t$$$ ($$$1 \leq t \leq 10^5$$$) — the number of test cases. The description of test cases follows.

The first line of each test case contains $$$n$$$ ($$$1 \le n \le 2 \cdot 10^5$$$) — the length of the string $$$s$$$.

The second line of each test case contains the binary string $$$s$$$ of length $$$n$$$.

It is guaranteed that sum of $$$n$$$ across all test cases doesn't exceed $$$2 \cdot 10^5$$$.

Output format:
For each test case, print a single integer — the number of good substrings.

Constraints:
Time limit: 2.0 s
Memory limit: 256.0 MB

Note:
Let's define a substring from index $$$l$$$ to index $$$r$$$ as $$$[l, r]$$$.

For the first test case, the good substrings are:

- $$$[1,1]$$$,
- $$$[1,2]$$$,
- $$$[3,6]$$$,
- $$$[4,5]$$$,
- $$$[4,6]$$$,
- $$$[5,5]$$$,
- $$$[5,6]$$$,
- $$$[6,6]$$$.

In the second test case, all substrings are good except $$$[2,2]$$$.

In the third test case, all substrings are good.

Example input:
4
6
100011
3
101
5
11111
6
010101

Example output:
8
5
15
18

Editorial:
We call a maximal sequence of $$$0$$$ a $$$0$$$ block, a maximal sequence of $$$1$$$ a $$$1$$$ block.
For each $$$i$$$, store some disjoint intervals with the following property:
For each $$$j$$$ that $$$i<=j$$$ and $$$s[j]=1$$$, j is in one of the intervals if and only if $$$(i,j)$$$ is a good substring.
We can prove the number of intervals for each $$$i$$$ is $$$O(\log n)$$$.
Let's start from $$$(i,i)$$$ to $$$(i,n)$$$, and consider when. good substring become bad and bad substring become good.
Assume we are at $$$(i,j)$$$.
If it is a good substring and ending with 1, it become bad when it meet a longer 0 block, so we get one interval and go to a bad substring ending with 0.
If it is a bad substring and ending with 0, suppose the next 1 is at k, then the next good intervals start at the smallest j' with three properties: (1) $$$(k,j')$$$ is good
(2) $$$s[j']=1$$$
(3) $$$(k,j')$$$ is not shorter than the last 0 block.
Proof: Property 2 is obvious because we only care about substrings ending with 1.
If $$$(k,j')$$$ is shorter than the last $$$0$$$ block, it is impossible to change this $$$0$$$ block, so $$$(i,j')$$$ is bad.
If $$$(k,j')$$$ is bad, if it starts with a 1 block $$$(k,j")$$$ not shorter than the last $$$0$$$ block after some operations, then $$$j"$$$ is smaller than $$$j'$$$ and has those three properties, otherwise it is impossible to change the last $$$0$$$ block.
If it has all three properties, in substring $$$(i,j')$$$, $$$(k,j')$$$ is good and can change the last $$$0$$$ block $$$1$$$, and then change $$$(i,j')$$$ to $$$1$$$, so we go to a good substring ending with 1.
When good substring become bad, the length is doubled, so the number of intervals for each i is $$$O(logn)$$$.
To know when good substring become bad ,for every 0 block $$$(j,k)$$$,if $$$(j,k)$$$ is longer than $$$(i,j-1)$$$, then $$$(i,k)$$$ is bad, we can preprocess those i in $$$O(n)$$$.
To know when bad substring become good, we go with $$$i$$$ from $$$n$$$ to $$$1$$$ and for each $$$i$$$, search in $$$O(\log n)$$$ intervals $$$O(\log n)$$$ times.
So we get those intervals in $$$O(n\log^2 n)$$$
Now we can check each substring ending with $$$1$$$ is good or not. Do the above algorithm on the reversed string, so we can check each substring starting with 1 is good or not. But we can not check substring starting and ending with 0.
Suppose $$$(i',j')$$$ is the longest $$$0$$$ block in $$$(i,j)$$$, then $$$(i,j)$$$ is good if at least one of $$$(i,i'-1)$$$ and $$$(j'+1,j)$$$ is good and not shorter than $$$(i',j')$$$.
Proof: After changing the longest $$$0$$$ block, we can change all other $$$0$$$ blocks. So for any substring, it is good if and only if we can change its longest $$$0$$$ block.
If $$$(i,i'-1)$$$ is good and not shorter, it can change $$$(i',j')$$$, so $$$(i,j)$$$ is good.
If $$$(i,j)$$$ is good assume the longest 0 block is changed from left. So after some operations in $$$(i,i'-1)$$$, there is a substring $$$(k,i'-1)$$$ become all $$$1$$$ and not shorter than $$$(i',j')$$$. Then $$$(k,i'-1)$$$ can make $$$(i,i'-1)$$$ good as there is no longer $$$0$$$ block.
Similarly, we can prove for $$$(j'+1,j)$$$.
So for each substring, we consider its longest $$$0$$$ block (substrings with no $$$0$$$ are always good). For convenience, if there are multiple maximum, we take the rightmost.
The longest $$$0$$$ block in substrings can be either a $$$0$$$ block in the original string or part of $$$0$$$ block in the original string. We can ignore substring in one $$$0$$$ block, so the longest $$$0$$$ block can only be a prefix or suffix, the number of possible longest $$$0$$$ block is $$$O(n)$$$.
For each possible 0 block $$$(i,j)$$$, assume the longest substring that $$$(i,j)$$$ is the longest $$$0$$$ block is $$$(i',j')$$$, we need to count answer for substring $$$(l,r)$$$ that $$$i' \le l \le i,j \le r \le j'$$$, it is equivalent to count the number of $$$l$$$ that $$$i'\le l \le i-(j-i+1)$$$ and $$$(l,i-1)$$$ is good and the number of $$$r$$$ that $$$j+(j-i+1) \le r \le j'$$$ and $$$(j+1,r)$$$ is good.
Notice that $$$s[i-1]=s[j+1]=1$$$, we can calculate them by above intervals.
We can use persistent segment tree directly or use segment tree or binary index tree after sorting queries. We need $$$O(nlogn)$$$ modifications and $$$O(n)$$$ queries. Thus we have solved our problem in $$$O(n\log^2 n)$$$.