Codeforces #906 : Div1 D. Game of Stacks

こっちの方がだいぶ実装が軽い。
https://codeforces.com/contest/1889/problem/D

問題

N個のスタックがあり、1～Nの番号が付いている。
各スタックには、1～Nのいずれかの整数が積まれている。

最初のスタック番号をpとしたとき、以下のアルゴリズムを考える。

• p番のスタックが空なら、pを返す。
• p番のスタックが空でないなら、p番のスタックの値を1個popし、pをその値にしたうえで、このアルゴリズムを繰り返す。

p=1～Nそれぞれに対し、このアルゴリズムが返す値を求めよ。

解法

スタックiに値A[i]がある場合、i→A[i]に辺を張った多重有向グラフを考える。
上記アルゴリズムは、最初点pに置いた駒を、辺に沿って動かす(その際通過した辺は削除する)時、pの終着点を求めるのに等しい。
注意点は、どの辺を選んでも良い点。元のアルゴリズムはスタックの上からpopする必要があるが、実際にはどの順に遷移しても結果が変わらない。

このグラフにおいて閉路を削除してしまおう。
あとはメモ化DPで、均しO(N)で各点からの終着点を求めることができる。

int N;
vector<int> E[202020];
int ret[202020];
int vis[202020];
vector<int> C;

int dfs(int cur) {
	/*
	cout<<cur<<" ";
	FORR(c,C) cout<<c;
	cout<<" ";
	cout<<E[0].size()<<E[1].size()<<E[2].size()<<endl;
	*/
	if(ret[cur]>=0) return ret[cur];
	if(E[cur].empty()) return ret[cur]=cur;
	if(vis[cur]) {
		while(C.back()!=cur) {
			E[C.back()].pop_back();
			vis[C.back()]=0;
			C.pop_back();
		}
		E[cur].pop_back();
		vis[cur]=0;
		C.pop_back();
		return dfs(cur);
	}
	vis[cur]=1;
	C.push_back(cur);
	return dfs(E[cur].back());
}

void solve() {
	int i,j,k,l,r,x,y; string s;
	
	cin>>N;
	FOR(i,N) {
		cin>>x;
		E[i].resize(x);
		FORR(e,E[i]) {
			cin>>e;
			e--;
		}
	}
	MINUS(ret);
	FOR(i,N) {
		C.clear();
		int tar=dfs(i);
		FORR(c,C) ret[c]=tar;
		cout<<ret[i]+1<<" ";
	}
	cout<<endl;
		
}

まとめ

順番が任意と気づくまでがしんどいね。