Codeforces Raif Round 1 : F. Fruit Sequences

まぁまぁ良い出来だった回。
https://codeforces.com/contest/1428/problem/F

問題

01で構成された文字列Sが与えられる。
f(L,R) := S[L...R]のうち、連続する1の最大数
とする。0≦L≦R＜|S|全通りに対するf(L,R)の総和を求めよ。

解法

BITを使い、以下をSを走査しながら更新する。
g(x) := 現在見ている位置をkとする。g(x)はS[g(x)...g(x)+x-1]が1をx文字並べたものと一致するような、最大の値を示す。
とすると、sum(f(x))がf(L,k)の総和となる。

int N;
string S;

template<class V, int ME> class BIT {
public:
	V bit[1<<ME],val[1<<ME];
	V operator()(int e) {if(e<0) return 0;V s=0;e++;while(e) s+=bit[e-1],e-=e&-e; return s;}
	void add(int e,V v) { val[e++]+=v; while(e<=1<<ME) bit[e-1]+=v,e+=e&-e;}
	void set(int e,V v) { add(e,v-val[e]);}
};
BIT<ll,20> bt;


void solve() {
	int i,j,k,l,r,x,y; string s;
	
	cin>>N>>S;
	ll ret=0;
	int len=0;
	for(i=1;i<=N;i++) {
		char c=S[i-1];
		if(c=='0') {
			for(j=1;j<=len;j++) bt.set(j,i-j);
			ret+=bt(N+1);
			len=0;
		}
		else {
			ret+=1LL*(i-len)*(len+1)+1LL*len*(len+1)/2;
			ret+=bt(N+1)-bt(len+1);
			len++;
		}
	}
	cout<<ret<<endl;
}

まとめ

本番は割とすんなり解いてるな。