yukicoder : No.2611 Count 01

最初の解法を思いついても、実装に手間取る問題。
https://yukicoder.me/problems/no/2611

問題

01で構成される文字列Xを考える。
f(X)とは、X中の0と1の頻度の積とする。
また、g(X)とは、Xの各部分文字列X'におけるf(X')の総和とする。

文字列Sが与えられるので、以下のクエリに順次答えよ。

• Sの1文字を0/1反転する
• Sの部分文字列S'が指定されるので、g(S')を答える。

解法

SegTreeで処理する。
各ノードには部分文字列S'に対応する以下を載せよう。

• 0/1の個数
• prefixにおける0/1の総和
• prefixにおける0と1の頻度の積の総和
• suffixにおける0/1の総和
• suffixにおける0と1の頻度の積の総和
• g(S')

const ll mo=998244353;
template<class V,int NV> class SegTree_1 {
public:
	vector<V> val;
	V def;
	V comp(V l,V r){
		V v;
		v[0]=l[0]+r[0];
		v[1]=l[1]+r[1];
		v[2]=l[2]+r[2]+r[0]*(l[0]+l[1]);
		v[3]=l[3]+r[3]+r[1]*(l[0]+l[1]);
		v[4]=l[4]+r[4]+r[0]*r[1]*(l[0]+l[1])+l[2]*r[1]+l[3]*r[0];
		v[5]=l[5]+r[5]+l[0]*(r[0]+r[1]);
		v[6]=l[6]+r[6]+l[1]*(r[0]+r[1]);
		v[7]=l[7]+r[7]+l[0]*l[1]*(r[0]+r[1])+r[5]*l[1]+r[6]*l[0];
		v[8]=l[8]+r[8]+l[4]*(r[0]+r[1])+r[7]*(l[0]+l[1])+l[2]*r[6]+l[3]*r[5];
		int i;
		FOR(i,9) v[i]=v[i]%mo;
		return v;
	};
	
	SegTree_1(){val=vector<V>(NV*2,def);};
	V getval(int x,int y,int l=0,int r=NV,int k=1) { // x<=i<y
		if(r<=x || y<=l) return def;
		if(x<=l && r<=y) return val[k];
		return comp(getval(x,y,l,(l+r)/2,k*2),getval(x,y,(l+r)/2,r,k*2+1));
	}
	void update(int entry, int v) {
		entry += NV;
		if(v==0) val[entry]={1,0,1,0,0,1,0,0,0};
		if(v==1) val[entry]={0,1,0,1,0,0,1,0,0};
		while(entry>1) entry>>=1, val[entry]=comp(val[entry*2],val[entry*2+1]);
	}
};
//0,1,a,b,ab,c,d,cd,g(x)
SegTree_1<array<ll,9>,1<<20> st;
int N,Q;

void solve() {
	int i,j,k,l,r,x,y; string s;
	
	cin>>N>>Q;
	cin>>s;
	FOR(i,N) {
		s[i]-='0';
		st.update(i,s[i]);
	}
	while(Q--) {
		cin>>i;
		if(i==1) {
			cin>>x;
			x--;
			s[x]^=1;
			st.update(x,s[x]);
		}
		else {
			cin>>x>>y;
			auto v=st.getval(x-1,y);
			cout<<v[8]<<endl;
		}
	}
}

まとめ

SegTreeでやるのはすぐ思いつくけど、計算が楽な値をうまくノードに載せるのに手間取る。