Codeforces #178 Div2. E. Shaass the Great

さてようやくE。
これ、自力で解けたは解けたけどだいぶ苦労した。
でもおかげで木構造のDFSが少し慣れてきた。
http://codeforces.com/contest/294/problem/E

問題

各辺にコストが振られた木構造のグラフが与えられる。
ここから１つ辺を取り除き、別の場所に同じコストの辺を追加してまた木構造を作ることを考える。
このとき、任意の２点間のコストの総和を最小化し、そのコストを答えよ。

解法

この問題は以下２つの手順が必要になる。

• 取り除く辺を選ぶこと
• 選んだ辺を付ける場所を選ぶこと

自分の解法では、前者は総当たりで全部の辺を試し、後者はもう少し効率よくやった。

まず、木構造から辺を１つ取り除くと２つの木L,Rに分かれる。
ここでL,Rをつなぐ辺をまたどこかにつなぐことになるが、その時のコストの総和を考える。
L中の点LaとR中の点Raを再度つなぐと仮定する。

• L内の2点間のコストは、辺La-Raを通らないので、La,Raの選択に限らず決まる
• R内の2点間のコストも同様。
• L内の点LbとR内の点Rb間のコストは、(Lb-La間のコスト)+(La-Ra間のコスト)+(Ra-Rb間のコスト)である。このとき、La-Raのコストは元々取り除いた辺のコストである。

これでやることは決まった。
Laは任意のL内の点とのコストの総和が最小化する点を選ぶと良い。Raも同様。
そのようなLaは、以下の手順でDFSを繰り返して求められる。

• まず、各辺についてその辺の両側にある点の総数をそれぞれ求める
• 適当なLaを定め、そこからL内の各点へのコストの総和を求める
• そこからDFSでLaの候補を移動して探していく。辺のコストと、辺の両端の点の総数がわかると辺を辿るたびにコストの変化分がわかるので、コストを最小化するLaを求める。

よって、全体でO(N^2)で終わる。
N<=5000なので、C++でも多くの人は１秒以上かかっているようだ。

int N;
vector<pair<int,int> > E[5002];
vector<ll> D[5002];
vector<int> Vs;
int CV[5002];
int CVP[5002];


void dfs(int cur, int pre, int dist) {
	int i;
	
	FOR(i,Vs.size()) D[max(cur,Vs[i])][min(cur,Vs[i])] = D[max(pre,Vs[i])][min(pre,Vs[i])] + dist;
	Vs.push_back(cur);
	
	FOR(i,E[cur].size()) {
		if(E[cur][i].first==pre) continue;
		dfs(E[cur][i].first,cur,E[cur][i].second);
	}
}

int numchild(int cur,int pre) {
	int i;
	if(CVP[cur]==pre) return CV[cur];
	CVP[cur]=pre;
	CV[cur]=1;
	FOR(i,E[cur].size()) {
		if(E[cur][i].first==pre) continue;
		CV[cur]+=numchild(E[cur][i].first, cur);
	}
	return CV[cur];
}

ll localtotal(int cur,int pre, ll tcnum) {
	int i;
	ll tt=0;
	
	FOR(i,E[cur].size()) {
		if(E[cur][i].first==pre) continue;
		tt += localtotal(E[cur][i].first, cur, tcnum) + 
			(ll)E[cur][i].second*(tcnum-CV[E[cur][i].first])*(ll)CV[E[cur][i].first];
	}
	return tt;
}

ll getmin(int cur,int pre, ll curt,ll tcnum) {
	ll mint=curt;
	int i;
	FOR(i,E[cur].size()) {
		if(E[cur][i].first==pre) continue;
		mint = min(mint, getmin(E[cur][i].first, cur, curt + E[cur][i].second*(tcnum-2*CV[E[cur][i].first]),tcnum));
	}
	return mint;
	
}

ll minroad(int X,int Y,ll C) {
	int i;
	ll txc=0,tyc=0,tt;
	
	numchild(X,Y);
	numchild(Y,X);
	tt = localtotal(X,Y,CV[X]) + localtotal(Y,X,CV[Y]);
	
	FOR(i,N) {
		if(D[max(i,X)][min(i,X)]<D[max(i,Y)][min(i,Y)]) txc += D[max(i,X)][min(i,X)];
		else tyc += D[max(i,Y)][min(i,Y)];
	}
	txc = getmin(X,Y,txc,CV[X]);
	tyc = getmin(Y,X,tyc,CV[Y]);
	
	return tt + txc*CV[Y]+tyc*CV[X]+C*CV[X]*(ll)CV[Y];
}

void solve() {
	int f,r,i,j,k,l,x,y,cur,tar;
	
	N=GETi();
	FOR(i,5000) D[i].resize(i+2,0);
	FOR(i,5000) CVP[i]=-2;
	Vs.reserve(5001);
	FOR(i,N-1) {
		x=GETi()-1;
		y=GETi()-1;
		r=GETi();
		E[x].push_back(make_pair(y,r));
		E[y].push_back(make_pair(x,r));
	}
	
	dfs(0,-1,0);
	
	ll mc=1LL<<60;
	FOR(x,N) FOR(y,E[x].size()) if(x<E[x][y].first) mc=min(mc, minroad(x,E[x][y].first,E[x][y].second));
	cout << mc << endl;
	
	return;
}

まとめ

何度もDFSして状態を更新する手法、この問題でだいぶ勉強になった。