Official Editorial (Python)

Solution - Prefix Sums + Math

Following the editorial, we build a prefix sum array $p$ on the existing array.

We know that subarray formed by $[l, r)$ is a good subarray iff $r-l=p_r-p_l$. Rearranging this equation leads to $p_r-r=p_l-l$, so we build a map (sum_dist) on values of $p_i-i$ for all valid $i$.

Then, we iterate over all values of the map and check how many unordered pairs we can build with the number of values of $i$ that have the same value of $p_i-i$.

Implementation

Time Complexity: $\mathcal{O}(N\log N)$

Copy
from collections import defaultdict

for _ in range(int(input())):
	arr_len = int(input())
	arr = [ord(i) - ord("0") for i in input()]

	pref_arr = [0] + arr
	for i in range(1, len(pref_arr)):
		pref_arr[i] += pref_arr[i - 1]