Codeforces 1428F Fruit Sequences 乱搞计数

Codeforces 1428F Fruit Sequences

题意

给一个长度为\(n\)的\(01\)串\(s\)，\(f(l,r)\)为区间\([l,r]\)内最长连续\(1\)的个数，计算\(\sum_{i=1}^{n}\sum_{j=i}^{n}f(i,j)\)。

\(n \le 5 \cdot 10^5\)

分析

设\(pre[i]\)为位置\(i\)之前连续\(1\)的个数，\(suf[i]\)为位置\(i\)之后连续\(1\)的个数，例如\(00111100\)，\(pre[4]=2,suf[4]=3\)。

令\(dp[i]\)为\(\sum_{j=1}^{i}f(j,i+suf[i]-1)\)，如何计算\(dp[i]\):

• 首先要加上\(\sum_{j=i-pre[i]+1}^{i}suf[j]\)，也就是包含位置\(i\)的一段连续\(1\)的贡献，可以用求和公式化简。
• 再找到\(i\)之前最近的一个位置\(j\)，且\(suf[j]>pre[i]+suf[i]-1\)，若找不到则令\(j=0\)，加上\(dp[j]+(i-pre[i]-j)\cdot (pre[i]+suf[i]-1)\)。

计算完\(dp\)数组后，再枚举端点\(r\)，令\(ans\)加上\(\sum_{j=1}^{pre[i]}j\)，也就是\(r\)之前连续一段\(1\)的贡献，再类似的找到\(r\)之前最近的一个位置\(j\)，且\(suf[j]>pre[r]\)，若找不到则令\(j=0\)，加上\(dp[j]+(r-pre[r]-j)\cdot pre[r]\)。

Code

#include <bits/stdc++.h>
#define lson l,mid,p<<1
#define rson mid+1,r,p<<1|1
using namespace std;
typedef long long ll;
const int N=5e5+10;
int n;
int a[N];
ll f[N];
int pre[N],suf[N],st[N][20];
void init(){
	for(int i=1;i<=n;i++) st[i][0]=suf[i];
	for(int j=1;(1<<j)<=n;j++){
		for(int i=1;i+(1<<j)-1<=n;i++){
			st[i][j]=max(st[i][j-1],st[i+(1<<(j-1))][j-1]);
		}
	}
}
int qmx(int l,int r){
	int k=(int)log2(r-l+1);
	return max(st[l][k],st[r-(1<<k)+1][k]);
}
int qy(int x,int num){
    int l=1,r=x;
    while(l<=r){
        int mid=l+r>>1;
        int cnt=qmx(mid,x);
        if(cnt>=num) l=mid+1;
        else r=mid-1;
    }
    return r;
}
int main()
{
    scanf("%d",&n);
    for(int i=1;i<=n;i++) scanf("%1d",&a[i]);

    for(int i=1;i<=n;i++){
        if(a[i]) pre[i]=pre[i-1]+1;
    }
    for(int i=n;i>=1;i--){
        if(a[i]) suf[i]=suf[i+1]+1;
    }
    init();
    for(int i=1;i<=n;i++) if(a[i]){
        f[i]=1ll*(2*suf[i]+pre[i]-1)*pre[i]/2;
        int x=pre[i]+suf[i]-1;
        int j=qy(i-pre[i],x+1);
        f[i]+=f[j];
        f[i]+=1ll*(i-pre[i]-j)*x;
    }
    ll ans=0;
    for(int i=1;i<=n;i++){
        ans+=1ll*pre[i]*(pre[i]+1)/2;
        int j=qy(i-pre[i],pre[i]+1);
        ans+=f[j];
        ans+=1ll*(i-pre[i]-j)*pre[i];
    }
    cout<<ans<<endl;
    return 0;
}