1622D - Shuffle - CodeForces Solution

#include<bits/stdc++.h>
#define mp make_pair
#define pii pair<int,int>
//#pragma GCC optimize ("O2")
#define pdd pair<double,double>
#define pll pair<ll,ll>
#define mst(a,b) memset(a,b,sizeof(a))
//#pragma comment(linker, "/STACK:102400000,102400000")
//g++ strlen.cpp -o a.exe -Wl,--stack,536870912
using namespace std;
typedef long long ll;
typedef unsigned long long ull;
typedef vector<double> vec;
typedef vector<vec> mat;
typedef complex <double> cp;
const double eps = 1e-8;
const int maxn = 5e3+9;
const double pi=acos(-1.0);
const ull mod1=19260817;
const ull mod2=19660813;
const int base=131;
const ll inf = 1e17;
const ll mod = 998244353;
ll op[maxn][maxn];
void init(int n,int k){
    op[0][0]=1;
    op[1][0]=op[1][1]=1;
    for(int i=2;i<=n;++i){
        op[i][0]=1;
        for(int j=1;j<=k&&j<=n;++j){
            op[i][j]=(op[i-1][j]+op[i-1][j-1])%mod;
        }
    }

}
char pq[maxn];
int main()
{
    ios::sync_with_stdio(false);
    int n,k;
    cin>>n>>k;
    init(n,k);
    cin>>pq;
    int l=0,r;

    if(k==0) cout<<1;
    else{
        while(l<n&&pq[l]!='1') ++l;
        ll ans=1;
        if(l==n) cout<<1;
        else{
            int tmp=0,p1=l,p2;
            l=0;
            for(r=p1;r<n;++r){
                if(pq[r]=='1'){
                    ++tmp;
                    if(tmp==k) p2=r;
                }
                if(tmp==k+1){
                    break;
                }
            }
//            cout<<p1<<","<<p2<<endl;
            if(tmp<k) cout<<1;
            else if(tmp==k) cout<<op[n][k];
            else{
                ans=(ans+op[r-l][k]-1+mod)%mod;
//                cout<<ans<<endl;
                while(r<n){
                    int l1,r1=r,p3=p1+1,p4;
                    p4=r;
                    ++r1;
                    while(r1<n&&pq[r1]=='0') ++r1;
                    l1=p1+1;
                    while(p3<n&&pq[p3]=='0') ++p3;
//                    cout<<"*********\n";
//                    cout<<l<<","<<p1<<","<<p2<<","<<r<<endl;
//                    cout<<l1<<","<<p3<<","<<p4<<","<<r1<<endl;
                    ans=(ans+op[r1-l1][k]-1+mod)%mod;
                    ans=(ans-op[p4-l1][k-1]+1+mod)%mod;
                    l=l1;r=r1;p1=p3;p2=p4;
//                    cout<<ans<<endl;
                }
                cout<<ans<<"\n";
            }
        }
    }
    return 0;
}