AtCoder ARC #128 (大和証券プログラミングコンテスト2021) : D - Neq Neq

苦戦しつつもなんとか通した。
https://atcoder.jp/contests/arc128/tasks/arc128_d

問題

N個の区別できるボールが並んでおり、それぞれ数値が振られている。
連続する3つのボールについて、真ん中のボールの値が両隣のボールと異なるなら、取り除くことができる。
最終的にあり得るボールの並びは何通りか。

解法

両端及び、同じ要素が連続している部分は消すことができない。
よって、それらの要素で元の列を分解し、要素隣接が互いに異なるような連続部分列に分けよう。
あとは連続部分列毎に組み合わせを求め、その積を取ろう。

以下は個々の連続部分列について考える。

f(x) := 先頭x個のボールまでを処理したとき、末尾をx番目のボールとするような並び順
とする。
f(1) = 1
f(y) = sum(f(x)) (ただしx＜yかつxとyの間のボールをすべて消せる場合)
となるので、f(y)を順に求めて行こう。

ほとんどの場合、xとyの間のボールはすべて消せるが、例外として両者の間に2つの値が交互になっているケースで、かつy-xが3以上の場合だけは取ることができない。
よってf(x)が確定したら、BITを使い上記に該当するy以外の部分に対し、f(y)にf(x)を足しこんでいこう。

int N;
int A[202020];
const ll mo=998244353;
int fix[202020];
int nex[202020];
int V[202020];


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

ll dp[202020];
int repeat[202020];

ll hoge(int L,int R) {
	if(L+1>=R) return 1;
	
	dp[L]=1;
	bt.add(L,1);
	bt.add(L+1,mo-1);
	for(int i=L;i<R;i++) {
		dp[i]=bt(i);
		
		if(repeat[i]>=2) {
			bt.add(i+1,dp[i]);
			bt.add(R+1,mo-dp[i]);
			bt.add(i+3,mo-dp[i]);
			bt.add(min(i+1+repeat[i]+1,R+1),dp[i]);
		}
		else {
			bt.add(i+1,dp[i]);
			bt.add(R+1,mo-dp[i]);
		}
	}
	//cout<<L<<R<<bt(R)<<endl;
	return bt(R);
}

void solve() {
	int i,j,k,l,r,x,y; string s;
	
	cin>>N;
	FOR(i,N) {
		cin>>A[i];
	}
	
	FOR(i,N+1) V[i]=N;
	for(i=N-1;i>=0;i--) {
		nex[i]=V[A[i]];
		V[A[i]]=i;
	}
	
	for(i=N-2;i>=0;i--) {
		if(A[i]==A[i+2]) {
			repeat[i]=repeat[i+1]+1;
		}
	}
	
	
	FOR(i,N) {
		if(i==0||i==N-1) fix[i]=1;
		else if(A[i]==A[i-1]||A[i]==A[i+1]) fix[i]=1;
	}
	int pre=-1;
	ll ret=1;
	FOR(i,N) {
		if(fix[i]==1) {
			if(pre!=-1) ret=ret*hoge(pre,i)%mo;
			pre=i;
		}
	}
	cout<<ret<<endl;
}

まとめ

これも何とか解けてよかった。