kmjp's blog

競技プログラミング参加記です

AtCoder ABC #246 : Ex - 01? Queries

いつものexよりは易しめ。
https://atcoder.jp/contests/abc246/tasks/abc246_h

問題

0/1/?で構成されるN文字の文字列Sが与えられる。
?の部分を0または1に置き換え、Sの部分列として構築できる空でない文字列の数をf(S)とする。
Sの1文字を置き換えるクエリが与えられるので、そのつどf(S)を答えよ。

解法

dp(n,0), dp(n,1)を、Sのprefix n文字の部分列として構築でき、末尾が0/1であるユニークな文字列の個数とする。

  • Sのn+1文字目が'0'の場合
    • dp(n+1,0)は、dp(n,0)にprefix n文字で作れる文字列dp(n,0)+dp(n,1)に加え'0'だけのものを追加で作れる。ただし、'0'を加えた結果dp(n,0)に含まれるものは二重カウントしている。よってdp(n+1,0) = dp(n,0)+dp(n,1)+1
    • dp(n+1,1)は、dp(n,1)に該当する文字列の末尾に'0'に加えても作れるものは変化しないのでdp(n+1,1)=dp(n,1)
  • Sのn+1文字目が'1'の場合
    • dp(n+1,0)=dp(n,0)
    • dp(n+1,1)=dp(n,0)+dp(n,1)+1
  • Sのn+1文字目が'?'の場合
    • dp(n+1,0)=dp(n+1,1)=dp(n,0)+dp(n,1)+1

上記処理より、3×3の行列の積でdp(N,*)を求めることができる。
3×3の行列をSegTreeで管理すれば、1点更新の累積積をO(logN)でとれる。

const int MAT=3;
struct Mat { ll v[MAT][MAT]; Mat(){ZERO(v);};};
const ll mo=998244353;

Mat mulmat(Mat& a,Mat& b,int n=MAT) {
	ll mo2=4*mo*mo;
	int x,y,z; Mat r;
	FOR(x,n) FOR(y,n) r.v[x][y]=0;
	FOR(x,n) FOR(z,n) FOR(y,n) {
		r.v[x][y] += a.v[x][z]*b.v[z][y];
		if(r.v[x][y]>mo2) r.v[x][y] -= mo2;
	}
	FOR(x,n) FOR(y,n) r.v[x][y]%=mo;
	return r;
}

template<class V,int NV> class SegTree_1 {
public:
	vector<V> val;
	V comp(V l,V r){ return mulmat(l,r);};
	
	SegTree_1(){
		val=vector<V>(NV*2);
		FORR(v,val) v.v[0][0]=v.v[1][1]=v.v[2][2]=1;
	};
	V getval(int x,int y,int l=0,int r=NV,int k=1) { // x<=i<y
		if(r<=x || y<=l) {
			Mat v;
			v.v[0][0]=v.v[1][1]=v.v[2][2]=1;
			return v;
		}
		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, V v) {
		entry += NV;
		val[entry]=v; //上書きかチェック
		while(entry>1) entry>>=1, val[entry]=comp(val[entry*2],val[entry*2+1]);
	}
};
SegTree_1<Mat,1<<20> st;
Mat A[3];
int N,Q;
string S;
void solve() {
	int i,j,k,l,r,x,y; string s;
	
	A[0].v[0][0]=A[0].v[0][1]=A[0].v[0][2]=A[0].v[1][1]=A[0].v[2][2]=1;
	A[1].v[0][0]=A[1].v[1][0]=A[1].v[1][1]=A[1].v[1][2]=A[1].v[2][2]=1;
	A[2].v[0][0]=A[2].v[0][1]=A[2].v[0][2]=A[2].v[1][0]=A[2].v[1][1]=A[2].v[1][2]=A[2].v[2][2]=1;
	
	cin>>N>>Q>>S;
	FOR(i,N) {
		if(S[i]=='0') st.update(i,A[0]);
		if(S[i]=='1') st.update(i,A[1]);
		if(S[i]=='?') st.update(i,A[2]);
	}
	FOR(i,Q) {
		cin>>x>>s;
		x--;
		if(s=="0") st.update(x,A[0]);
		if(s=="1") st.update(x,A[1]);
		if(s=="?") st.update(x,A[2]);
		cout<<(st.val[1].v[0][2]+st.val[1].v[1][2])%mo<<endl;
	}
	
}

まとめ

とはいえ遷移時の組み合わせ計算が若干トリッキー。