PROBLEM LINK:
Practice
Contest: Division 1
Contest: Division 2
Contest: Division 3
Contest: Division 4
Author: raysh07
Tester: apoorv_me
Editorialist: iceknight1093
DIFFICULTY:
TBD
PREREQUISITES:
Prefix sums
PROBLEM:
You’re given a binary string S.
Find the sum of the number of modes across all substrings of S.
EXPLANATION:
S contains only the characters 0 and 1, so for any substring of S, exactly one of these three possibilities will hold:
- There is a unique mode, that being 0.
- There is a unique mode, that being 1.
- Both 0 and 1 are modes of the substring.
In particular, notice that the third condition — both 0 and 1 being modes — can happen if and only if the number of zeros in the substring equals the number of ones (after all, if there were more zeros than ones then zero would be the unique mode).
So, every substring of S that has an equal number of zeros and ones has 2 modes, everything else has 1 mode.
Let’s try to count the number of substrings of S that have an equal number of zeros and ones.
To do this, we use a small trick: we look at the difference between the number of zeros and ones, and try to look at when this is 0.
Specifically, consider an array A of length N, where A_i = 1 if S_i = 1, and A_i = -1 if S_i = 0.
For example, if S = \texttt{10010} we’ll get A = [1, -1, -1, 1, -1].
Then, observe that for any substring S[l\ldots r] of S,
- If S[l\ldots r] contains strictly more zeros than ones, A_l + A_{l+1} + \ldots + A_r \lt 0.
The converse holds too. - If S[l\ldots r] contains strictly more ones than zeros, A_l + A_{l+1} + \ldots + A_r \gt 0, and vice versa.
- If S[l\ldots r] contains an equal number of zeros and ones, A_l + A_{l+1} + \ldots + A_r = 0 will hold.
That is, all we really want to do is count the number of subarrays of A whose sum is zero!
This is a relatively standard task that can be solved using prefix sums.
How?
Let P denote the prefix sum array of A, so that P_0 = 0 and P_i = P_{i-1} + A_i for i\geq 1.
Then, the sum of the subarray [l, r] of A is P_r - P_{l-1}, and of course this will equal 0 exactly when P_r = P_{l-1}.
So, counting the number of subarrays whose sum is 0 is equivalent to counting the number of pairs of equal elements in P.
The latter is doable with some simple combinatorics: if the value x appears f_x times in P, then it contributes \frac{f_x\cdot (f_x - 1)}{2} pairs of subarrays with sum 0, so the problem can be solved in linear time by just computing the frequency table of P.
Once you know the number of subarrays with sum 0 (equivalently, substrings with equal zeros and ones), they contribute 2 each to the answer while everything else contributes 1 (the count of the others is easily obtained, since there are \frac{N\cdot (N+1)}{2} substrings in total).
TIME COMPLEXITY:
\mathcal{O}(N) per testcase.
CODE:
Author's code (C++)
#include <bits/stdc++.h>
using namespace std;
#define int long long
#define INF (int)1e18
#define f first
#define s second
mt19937_64 RNG(chrono::steady_clock::now().time_since_epoch().count());
void Solve()
{
int n; cin >> n;
string s; cin >> s;
int curr = 0;
map <int, int> v;
v[0]++;
int ans = n * (n + 1) / 2;
for (auto x : s){
if (x == '1') curr++;
else curr--;
ans += v[curr]++;
}
cout << ans << "\n";
}
int32_t main()
{
auto begin = std::chrono::high_resolution_clock::now();
ios_base::sync_with_stdio(0);
cin.tie(0);
int t = 1;
// freopen("in", "r", stdin);
// freopen("out", "w", stdout);
cin >> t;
for(int i = 1; i <= t; i++)
{
//cout << "Case #" << i << ": ";
Solve();
}
auto end = std::chrono::high_resolution_clock::now();
auto elapsed = std::chrono::duration_cast<std::chrono::nanoseconds>(end - begin);
cerr << "Time measured: " << elapsed.count() * 1e-9 << " seconds.\n";
return 0;
}
Editorialist's code (Python)
for _ in range(int(input())):
n = int(input())
s = input()
ans = pref = 0
from collections import defaultdict
ct = defaultdict(int)
ct[0] = 1
for c in s:
pref += 1 if c == '1' else -1
ans += ct[pref]
ct[pref] += 1
print(ans + n*(n+1)//2)