Bit strings

#include <bits/stdc++.h>
#include <cmath>
#include <complex>
using namespace std;

const int N = 1e5 + 5;

const double PI = acos(-1.0);
struct Comp {
    double x, y;
    Comp(double _x=0.0, double _y=0.0)
    {
        x = _x;
        y = _y;
    }
    Comp operator - (const Comp &P)const {
        return Comp(x-P.x, y-P.y);
    }
    Comp operator + (const Comp &P)const {
        return Comp(x+P.x, y+P.y);
    }
    Comp operator * (const Comp &P)const {
        return Comp(x*P.x-y*P.y, x*P.y+y*P.x);
    }
};

void Change(Comp y[], int len)
{
    int i, j, k;
    for(i = 1, j=len/2; i < len - 1; i++)
    {
        if(i < j) swap(y[i], y[j]);
        k = len/2;
        while(j >= k)
        {
            j -= k;
            k /= 2;
        }
        if(j < k) j += k;
    }
}
void FFT(Comp y[], int len, int on)
{
    // on = 1 for DFT, on = -1 for IDFT
    Change(y, len);
    for(int h = 2; h<=len; h <<= 1)
    {
        Comp wn(cos(-on*2*PI/h), sin(-on*2*PI/h));
        for(int j=0; j<len; j+=h)
        {
            Comp w(1, 0);
            for(int k = j; k < j+h/2; k++)
            {
                Comp u = y[k];
                Comp t = w*y[k+h/2];
                y[k] = u+t;
                y[k+h/2] = u-t;
                w = w*wn;
            }
        }
    }
    if(on == -1) for(int i=0; i<len; i++)
        y[i].x /= len;
}

Comp x1[4*N], x2[4*N];
int l[N], z[N], a[N];
int c[N];
void solve()
{
    string s;
    cin >> s;

    int m = s.length(), k=0;
    memset(l, 0, sizeof(int)*(m+5));
    memset(z, 0, sizeof(int)*(m+5));
    memset(a, 0, sizeof(int)*(m+5));

    int last = -1;
    for(int i=0; i<m; i++)
    {
        z[i] = z[i-1];
        if(s[i]=='1')
        {
            // in one
            k++;
            l[k] = z[i] - z[last] + 1;
            last = i;
        } else {
            // in zero
            z[i]+=1;
        }
    }
    l[k+1] = z[m-1] - z[last] + 1;

    // for(int i=1; i<=k; i++)
    // {
    //     a[0] += c[l[i] - 1];
    //     for(int j=0; i+j<=k; j++)
    //         a[j+1] += l[i] * l[i+j+1];
    // }
    // a[0] += c[l[k+1] - 1];

    // New solution
    for(int i=1; i<=k+1; i++)
        a[0] += c[l[i] - 1];

    int len = 1;
    while(len < 2*k)
        len <<= 1;

    for(int i=0; i<k; i++)  
    {
        x1[i] = Comp(l[i+1], 0);
        x2[i] = Comp(l[k-i+1], 0);
    }
    for(int i=k; i<len; i++)
    {
        x1[i] = Comp(0, 0);
        x2[i] = Comp(0, 0);
    }

// for(int i=0; i<len; i++)
//     cout<<x1[i].x<<" ";
// cout<<endl;
// for(int i=0; i<len; i++)
//     cout<<x2[i].x<<" ";
// cout<<endl;

    FFT(x1, len, 1);
    FFT(x2, len, 1);
    for(int i=0; i<len; i++)
        x1[i] = x1[i] * x2[i];
    FFT(x1, len, -1);
// for(int i=0; i<len; i++)
//     cout<<x1[i].x<<" ";
// cout<<endl;
    for(int i=0; i<k; i++)
        a[k-i] = (int)(x1[i].x+0.5);

    for(int i=0; i<=m; i++)
    {
        // all numbers
        cout<<a[i]<<" ";
    }
    cout<<endl;
}


signed main() {
    cin.tie(0)->sync_with_stdio(0);

    int t;
    cin >> t;

    c[1] = 1;
    for(int i=2; i<=1e5; i++)
        c[i] = c[i-1] + i;
    
    while (t--) {
        solve();
    }
}

Test 1

Verdict: ACCEPTED

Test 2

Verdict: ACCEPTED

Test 3

Verdict: ACCEPTED

Test 4

Verdict: ACCEPTED

Test 5

Verdict: ACCEPTED

Test 6

Verdict: ACCEPTED