PROBLEM LINK:
Setter: Anik Sarker, Ezio Auditore
Tester: Teja Vardhan Reddy
Editorialist: Taranpreet Singh
DIFFICULTY:
PREREQUISITES:
Prefix Sums, Sliding Window, brute-force, amortization.
PROBLEM:
Given a string S of length N consisting of characters ‘0’ and ‘1’ only. Find the number of non-empty substrings of S such that cnt_0 = cnt_1*cnt_1 where cnt_c denote the frequency of character c in substring.
EXPLANATION
Let us assume cnt_1 = x for some substring. Then cnt_0 = x^2, hence the length of substring should be cnt_0+cnt_1 = x^2+x since we cannot have any other character in substring. Hence, we can surely reject all substrings whose length cannot be written as x^2+x for some integer x.
Now, we can see, the maximum value of x such that x^2+x \leq N is around \sqrt{N}. So, what we can do is to try all possible values of x and for each x, consider all substrings of S with length x^2+x and out of these, count the number of substrings which contain exactly x occurrences of character ‘1’. We can see, that this substring shall contain exactly x^2+x-x = x^2 occurrences of character ‘0’ then, which satisfies our requirement.
Code Snippet
for(int x = 1; x*x+x <= N; x++)
//we need to consider all substrings of length (x^2+x) and count the substrings, which have x occurrences of 1
For doing this, we have two ways, sliding window, and prefix-sums. In the sliding window, we can make a counter variable counting the number of occurrences of character ‘1’ in the current range. Both of these allows us to consider all substrings of a specific length in O(N) time.
It can be seen, that we need to consider \sqrt{N} different lengths of substrings and considering substrings of each length takes O(N) time, we have solved this problem in O(N*\sqrt{N}) time.
TIME COMPLEXITY
The time complexity is O(N*\sqrt{N}) per test case.
SOLUTIONS:
Setter 1 Solution
#include<bits/stdc++.h>
using namespace std;
#define MAX 100005
int C[MAX];
int main(){
int t;
scanf("%d",&t);
assert(1 <= t && t<= 10);
for(int cs=1;cs<=t;cs++){
string s;
cin>>s;
int n = s.size();
assert(1 <= n && n<= 100000);
for(int i=1;i<=n;i++){
assert(s[i-1] == '0' || s[i-1] == '1');
C[i] = C[i-1] + s[i-1] - '0';
}
int Count = 0;
for(int i=0;i<n;i++){
int x = 1;
while(true){
int j = i + x*x + x;
if(j > n) break;
int One = C[j] - C[i];
if(One == x) Count++;
x++;
}
}
printf("%d\n",Count);
}
}
Setter 2 Solution
#include<bits/stdc++.h>
using namespace std;
typedef long long ll;
const int maxn = 100009;
int n, m, x, y, z;
string str;
int csum[maxn];
int main()
{
// freopen("input05.txt", "r", stdin);
// freopen("output05.txt", "w", stdout);
int t, cs = 1;
cin >> t;
if(t < 1 || t > 10) assert(false);
while(t--){
cin >> str;
memset(csum, 0, sizeof(csum));
if(str.size() > 100000) assert(false);
int n = str.size();
str = " " + str;
for(int i = 1; i <= n; i++){
if(str[i] == '0') csum[i] = csum[i - 1];
else if(str[i] == '1') csum[i] = csum[i - 1] + 1;
else{
assert(false);
}
}
int ans = 0;
for(int i = 1; ; i++){
int len = i * i + i;
if(len > n) break;
for(int j = 1; j <= n - len + 1; j++){
long long cnt1 = csum[j + len - 1] - csum[j - 1];
long long cnt0 = len - cnt1;
if(cnt0 == cnt1 * cnt1) ans++;
}
}
printf("%d\n", ans);
}
return 0;
}
Tester's Solution
//teja349
#include <bits/stdc++.h>
#include <vector>
#include <set>
#include <map>
#include <string>
#include <cstdio>
#include <cstdlib>
#include <climits>
#include <utility>
#include <algorithm>
#include <cmath>
#include <queue>
#include <stack>
#include <iomanip>
#include <ext/pb_ds/assoc_container.hpp>
#include <ext/pb_ds/tree_policy.hpp>
//setbase - cout << setbase (16); cout << 100 << endl; Prints 64
//setfill - cout << setfill ('x') << setw (5); cout << 77 << endl; prints xxx77
//setprecision - cout << setprecision (14) << f << endl; Prints x.xxxx
//cout.precision(x) cout<<fixed<<val; // prints x digits after decimal in val
using namespace std;
using namespace __gnu_pbds;
#define f(i,a,b) for(i=a;i<b;i++)
#define rep(i,n) f(i,0,n)
#define fd(i,a,b) for(i=a;i>=b;i--)
#define pb push_back
#define mp make_pair
#define vi vector< int >
#define vl vector< ll >
#define ss second
#define ff first
#define ll long long
#define pii pair< int,int >
#define pll pair< ll,ll >
#define sz(a) a.size()
#define inf (1000*1000*1000+5)
#define all(a) a.begin(),a.end()
#define tri pair<int,pii>
#define vii vector<pii>
#define vll vector<pll>
#define viii vector<tri>
#define mod (1000*1000*1000+7)
#define pqueue priority_queue< int >
#define pdqueue priority_queue< int,vi ,greater< int > >
#define flush fflush(stdout)
#define primeDEN 727999983
mt19937 rng(chrono::steady_clock::now().time_since_epoch().count());
// find_by_order() // order_of_key
typedef tree<
int,
null_type,
less<int>,
rb_tree_tag,
tree_order_statistics_node_update>
ordered_set;
#define int ll
int a[123456];
main(){
std::ios::sync_with_stdio(false); cin.tie(NULL);
int t;
cin>>t;
while(t--){
int i;
string s;
cin>>s;
int cnt=0;
int ans=0;
//map<int,int> mapi;
int n=s.length(),j;
rep(i,s.length()){
if(s[i]=='1')
cnt++;
a[i+1]=cnt;
}
//int ans=0;
rep(i,s.length()){
//previ=a[i];
for(j=1;j*j+j+i<=n;j++){
if(a[j*j+j+i]-a[i]==j){
ans++;
}
}
}
cout<<ans<<endl;
}
return 0;
}
Editorialist's Solution
import java.util.*;
import java.io.*;
import java.text.*;
class BDGFT{
//SOLUTION BEGIN
//Into the Hardware Mode
void pre() throws Exception{}
void solve(int TC) throws Exception{
String s = n();int n = s.length();
int[] pre = new int[1+n];
for(int i = 1; i<= n; i++){
pre[i] = pre[i-1];
if(s.charAt(i-1) == '1')pre[i]++;
}
long ans = 0;
for(int i = 1; i<= n; i++)
//Iterating over all values of cnt1 till it exceed length of substring S[i, n]
for(int cnt1 = 1; i-1+cnt1+cnt1*cnt1 <= n; cnt1++)
//Checking if substring of length cnt1+cnt1*cnt1 has cnt1 ones.
if(pre[i-1+cnt1+cnt1*cnt1]-pre[i-1] == cnt1)
ans++;
pn(ans);
}
//SOLUTION END
void hold(boolean b)throws Exception{if(!b)throw new Exception("Hold right there, Sparky!");}
long IINF = (long)1e18, mod = (long)1e9+7;
final int INF = (int)1e9, MX = (int)2e5+5;
DecimalFormat df = new DecimalFormat("0.00000000000");
double PI = 3.141592653589793238462643383279502884197169399, eps = 1e-6;
static boolean multipleTC = true, memory = false, fileIO = false;
FastReader in;PrintWriter out;
void run() throws Exception{
if(fileIO){
in = new FastReader("input.txt");
out = new PrintWriter("output.txt");
}else {
in = new FastReader();
out = new PrintWriter(System.out);
}
//Solution Credits: Taranpreet Singh
int T = (multipleTC)?ni():1;
pre();for(int t = 1; t<= T; t++)solve(t);
out.flush();
out.close();
}
public static void main(String[] args) throws Exception{
if(memory)new Thread(null, new Runnable() {public void run(){try{new BDGFT().run();}catch(Exception e){e.printStackTrace();}}}, "1", 1 << 28).start();
else new BDGFT().run();
}
long gcd(long a, long b){return (b==0)?a:gcd(b,a%b);}
int gcd(int a, int b){return (b==0)?a:gcd(b,a%b);}
int bit(long n){return (n==0)?0:(1+bit(n&(n-1)));}
void p(Object o){out.print(o);}
void pn(Object o){out.println(o);}
void pni(Object o){out.println(o);out.flush();}
String n()throws Exception{return in.next();}
String nln()throws Exception{return in.nextLine();}
int ni()throws Exception{return Integer.parseInt(in.next());}
long nl()throws Exception{return Long.parseLong(in.next());}
double nd()throws Exception{return Double.parseDouble(in.next());}
class FastReader{
BufferedReader br;
StringTokenizer st;
public FastReader(){
br = new BufferedReader(new InputStreamReader(System.in));
}
public FastReader(String s) throws Exception{
br = new BufferedReader(new FileReader(s));
}
String next() throws Exception{
while (st == null || !st.hasMoreElements()){
try{
st = new StringTokenizer(br.readLine());
}catch (IOException e){
throw new Exception(e.toString());
}
}
return st.nextToken();
}
String nextLine() throws Exception{
String str = "";
try{
str = br.readLine();
}catch (IOException e){
throw new Exception(e.toString());
}
return str;
}
}
}
Feel free to share your approach, if you want to. (even if its same 
) . Suggestions are welcomed as always had been. 
