Codeforces
CF Step
Youtube Linkedin Discord Toggle Dark/Light/Auto mode Toggle Dark/Light/Auto mode Toggle Dark/Light/Auto mode Back to homepage

Hints: Decode

Section 1: Whole array good pairs (fix $L=1, R=n$)

Warm-up: set $L=1$ and $R=n$. Then the task is to count how many pairs $(x,y)$ with $1 \le x \le y \le n$ make the substring $s_x..s_y$ have the same number of 0 and 1. How can “equal counts” be rewritten as an equation?

Map the string to $\pm 1$ values so that equal counts becomes sum equals zero.

Let $$ a_i=\begin{cases} +1 & s_i=\mathtt{0}\ -1 & s_i=\mathtt{1} \end{cases} $$ Then for a substring $[x,y]$, the numbers of 0 and 1 are equal iff $$ \sum_{i=x}^{y} a_i = 0. $$

Answer to Hint 2: Use prefix sums $p_0=0$ and $p_i=\sum_{k=1}^{i} a_k$.

Then $$ \sum_{i=x}^{y} a_i = p_y - p_{x-1}, $$ so the substring is good iff $p_y=p_{x-1}$.

For the warm-up (whole array), every good pair $(x,y)$ corresponds to choosing two prefix indices $(x-1,y)$ with the same prefix value. If a value $v$ occurs $cnt[v]$ times among $p_0..p_n$, the number of good pairs is: $$ \sum_v \binom{cnt[v]}{2}. $$

Section 2: All subarrays

Returning to the original question: for every subarray $[l,r]$, count good pairs $(x,y)$ with $l \le x \le y \le r$, and sum over all $(l,r)$.

Instead of recomputing per $[l,r]$, flip the summation: for a fixed good pair $(x,y)$, count how many subarrays $[l,r]$ include it.

  • $l$ can be any value in $[1,x]$ so there are $x$ choices,
  • $r$ can be any value in $[y,n]$ so there are $(n-y+1)$ choices.

So $$ \text{ans}=\sum_{1\le x\le y\le n} x,(n-y+1)\cdot [p_{x-1}=p_y]. $$

Now fix the right endpoint $y$. The condition $p_y=p_{x-1}$ only depends on the value of the prefix sum.

Rewrite the weighted sum from Hint 4: $$ \sum_{x=1}^{y} x,(n-y+1)\cdot [p_{x-1}=p_y]. $$ If you set $j=x-1$, then $x=j+1$ and $j<y$.

Answer to Hint 5: Maintain for each prefix value $v$ the accumulated quantity $$ f[v]=\sum_{\substack{j\text{ seen so far}\ p_j=v}} (j+1). $$

When you are at position $y$, the contribution becomes: $$ (n-y+1)\cdot f[p_y]. $$

After adding the contribution for $y$, update so this $y$ can be used as a future $j$.

Since $j=y$ has weight $(j+1)=(y+1)$, do: $$ f[p_y] \mathrel{+}= (y+1). $$

Also initialize the “empty prefix” case: $p_0=0$, so start with $f[0]=1$ (because $0+1=1$).

Implementation tips:

  • The prefix sums $p_i$ range from $-n$ to $n$, so you can store $f$ in an array indexed by $(p_i+n)$, or use a hashmap.
  • Iterate $y=1..n$ once per test case: total complexity $O(n)$ per string.
  • Take everything modulo $10^9+7$.