kmjp's blog

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

Codeforces #606 Div1 D. Tree Elimination

一転こちらは割と難しめ。
https://codeforces.com/contest/1276/problem/D

問題

木を成す無向グラフが与えられる。
辺には番号が振られる。また、初期状態で各頂点にトークンが置いてある。

辺を番号順に見ていくとき、両端の頂点のいずれにもトークンがあれば、それを取り除く、ということを考える。
取り除いたトークンの頂点番号を並べると何通りあるか。

解法

各頂点vとその親pについてみたとき、vのパターンとしては

  • p-vの辺を処理する前に、v-(vの子頂点)の辺を処理した際にvを取り除く
  • p-vの辺を処理したときpを取り除く
  • p-vの辺を処理したときvを取り除く

を考える。(p-vの前にpを取り除くパターンは、pにおける最初のパターンで考慮する)。

それぞれのケースにおいてDPで数え上げて行こう。
式の遷移はEditorialを参照して頂くとしてここには書かないが、累積積とか出てきて結構遷移が面倒なものの、O(|V|)で済む。

int N;
vector<pair<int,int>> E[202020];
int U[202020],V[202020];
const ll mo=998244353;

ll dp[202020][3];

void dfs(int cur,int pre) {
	sort(ALL(E[cur]));
	int d=0;
	vector<ll> prefix,suffix;
	vector<ll> vs;
	prefix.push_back(1);
	suffix.push_back(1);
	FORR(e,E[cur]) {
		if(e.second==pre) {
			d=lower_bound(ALL(E[cur]),e)-E[cur].begin();
		}
		else {
			dfs(e.second,cur);
			prefix.push_back(prefix.back()*(dp[e.second][0]+dp[e.second][1])%mo);
			vs.push_back(e.second);
		}
	}
	reverse(ALL(E[cur]));
	FORR(e,E[cur]) if(e.second!=pre) {
		suffix.push_back(suffix.back()*(dp[e.second][0]+dp[e.second][2])%mo);
	}
	reverse(ALL(E[cur]));
	reverse(ALL(suffix));
	int i;
	dp[cur][1]=prefix[d]*suffix[d]%mo;
	dp[cur][2]=prefix.back();
	FOR(i,vs.size()) {
		if(i<d) (dp[cur][0]+=prefix[i]*dp[vs[i]][2]%mo*suffix[i+1])%=mo;
		else (dp[cur][2]+=prefix[i]*dp[vs[i]][2]%mo*suffix[i+1])%=mo;
	}
	
}


void solve() {
	int i,j,k,l,r,x,y; string s;
	
	cin>>N;
	FOR(i,N-1) {
		cin>>U[i]>>V[i];
		U[i]--;
		V[i]--;
		E[U[i]].push_back({i,V[i]});
		E[V[i]].push_back({i,U[i]});
	}
	E[0].push_back({N-1,N});
	
	dfs(0,N);
	cout<<(dp[0][0]+dp[0][1])%mo<<endl;
}

まとめ

これ解いた時まだWindows7だもんな…。