Codeforces #591 Div1 D. Stack Exterminable Arrays

Div1Dの割には簡単なのだが、自分は本番ミスしてた。
https://codeforces.com/contest/1240/problem/D

問題

数列Aに対し、以下の手順を行う。

• 数列の要素を先頭からスタックに積むことを考える。
  • 積もうとしたとき、すでにスタックの最上位に同じ値があれば、それを消し、積むのもやめる。
  • そうでなければスタックの最上位に積む。

Aの部分列に対し、上記手順で消しきれるものは何通りか。

解法

以下をmapで持つ。
f(i,p) := A[i]を先頭とする数列で、A[i...f(i,p)-1]までちょうど消しきることができ、A[f(i,p)]=A[i]となる最小の位置
また、
dp(i) := A[i]を先頭とする部分列のうち、全体を消しきれるものの数
とする。上記値をiの大きい順に求めていこう。

A[i]を先頭とするときどうなるかというと、A[i...f(i+1,A[i])]までが最初に消える位置になるので、
dp(i) = dp(f(i+1,A[i]))+1
となる。またその時

• f(i,*) = f(f(i+1,A[i]),*)
• ただしf(i,A[i])=i

となる。

int Q;
int N;
int A[303030];
int R[303030];
map<int,int> M[303030];
ll dp[303030];
void solve() {
	int i,j,k,l,r,x,y; string s;
	
	srand(time(NULL));
	cin>>Q;
	while(Q--) {
		cin>>N;
		
		FOR(i,N) {
			cin>>A[i];
			M[i].clear();
		}
		M[N].clear();
		dp[N]=0;
		
		ll ret=0;
		for(i=N-1;i>=0;i--) {
			if(M[i+1].count(A[i])) {
				R[i]=M[i+1][A[i]];
				swap(M[i],M[R[i]+1]);
				dp[i]=1+dp[R[i]+1];
			}
			else {
				dp[i]=0;
				R[i]=N;
			}
			M[i][A[i]]=i;
			ret+=dp[i];
		}
		
		cout<<ret<<endl;
	}
}

まとめ

本番は考察漏れで落としてた。残念。