kmjp's blog

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

Codeforces #277 Div2 D. Valid Sets

ちょっと迷ったけど最終的には解けて良かった。
http://codeforces.com/contest/351/problem/D

問題

N頂点と無向辺からなる木を成すグラフが与えられる。
各頂点には整数値A[i]が降られている。

ここで、頂点群の部分集合Sのうち、以下を満たすものの組み合わせ数を答えよ。

  • Sは空でない。
  • Sを構成する頂点群は隣接している。
  • Sに含まれる頂点群のA[i]の最大値と最小値の差はD以下である。

解法

まずは自分の解法を紹介。

まず木を根付き木とみなす。
自分は状態として[頂点番号][Sの最小値]を持ち、その番号の頂点のサブツリーでの組み合わせとして(Sの最小値を含む連結頂点群の組み合わせ数, Sの最小値を含まない連結頂点群の組み合わせ数)の2値を計算した。
この時の連結頂点群は、(Sの最小値)≦A[i]≦(Sの最小値)+D以下のものしか含むことができない。
先にSの最小値と、頂点群の最も根に近い点を固定することで一意性を保つ。

サブツリーの[頂点番号][Sの最小値]の値の計算では、Sの最小値を含む連結頂点群の組み合わせ数は、各子頂点における(Sの最小値を含む連結頂点群の組み合わせ + Sの最小値を含まない連結頂点群の組み合わせ)の積から、(Sの最小値を含まない連結頂点群の組み合わせ)の積を引くことで、最低1個はSの最小値を含む組み合わせを得ることができる。

あとは[頂点番号][Sの最小値]の各状態を求め、「Sの最小値を含むSの最小値を含む連結頂点群の組み合わせ数」の総和を求めればよい。


なお、自分のこの解法は若干長くてややこしい。
他にはもっと短い解法もある。
それは各頂点iのA[i]を最小値とするSの組み合わせ数をDFSで求める方法である。
頂点jはA[i]<A[j]≦A[i]+D またはA[i]==A[j]かつi<jな場合のみ連結可能とする。
後者の頂点番号の大小関係を用いることで、同じ最小値が複数回登場するSでも、多重にカウントすることを防いでいる。

Codeforcesで見られる短い解法はこちらを使ってるものが多いようだ。

int D,N;
int A[2001];
vector<int> E[2001];
pair<ll,ll> dp[2002][2002];
ll mo=1000000007;

pair<ll,ll> dodo(int cur,int mi) {
	if(dp[cur][mi].first>=0) return dp[cur][mi];
	if(A[cur]>mi+D || A[cur]<mi) return make_pair(0,1);
	
	int i;
	ll yes=1,no=1;
	FOR(i,E[cur].size()) {
		pair<ll,ll> p=dodo(E[cur][i],mi);
		yes = yes*(p.first+p.second)%mo;
		no = no*p.second%mo;
	}
	
	if(A[cur]==mi) no = 1;
	else yes = (((yes-no)%mo)+mo)%mo, no=(no+1)%mo;
	return dp[cur][mi]=make_pair(yes,no);
}

int dfs(int cur,int pre) {
	int i,p=-1;
	FOR(i,E[cur].size()) {
		if(E[cur][i]==pre) p=i;
		else dfs(E[cur][i],cur);
	}
	if(p>=0) E[cur].erase(E[cur].begin()+p);
}

void solve() {
	int i,j,k,l,r,x,y; string s;
	
	cin>>D>>N;
	FOR(i,N) cin>>A[i];
	FOR(i,N-1) {
		cin>>x>>y;
		E[x-1].push_back(y-1);
		E[y-1].push_back(x-1);
	}
	dfs(0,-1);
	MINUS(dp);
	ll ret=0;
	for(i=1;i<=2000;i++) FOR(x,N) {
		pair<ll,ll> p=dodo(x,i);
		ret=(ret+p.first)%mo;
	}
	cout<<ret<<endl;
}

まとめ

ちょっと手間をかけすぎたけど、まぁ解けたのでよいか。