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;
}