Codeforces #411 Div1 D. Expected diameter of a tree

計算量の見積もりが甘かったけど何とか通った。
http://codeforces.com/contest/804/problem/D

問題

森を成す無向グラフが与えられる。
この森に対し、以下のクエリに答えよ。

クエリは２頂点(u,v)からなる。
uの連結成分中の頂点a、vの連結成分中の頂点bを均等な確率で選び、その間を辺で結ぶ。
そうして２つの木を連結した部分グラフの直径の期待値を求めよ。

解法

まず個々の木に対して直径を求め、そこからさらに各頂点から最遠点までの距離を求めておこう。
次に、各連結成分について最遠点の距離を昇順ソートしたものと、その累積和を求めておく。

各クエリについて、U,Vをそれぞれu,vの連結成分における最遠点の距離を昇順ソートしたものとする。
辺で連結することで直径が伸びることもあれば、元の両連結成分の直径の方がもともと長い可能性もある。
よって求める解はとなる。

上記計算をまともにやると、最悪O(N^2*Q)かかりTLEするため、以下の２つのテクで計算量を落とした。

• 結果をキャッシュして同じ計算を繰り返さないようにする
  • 頂点数の多い連結成分の数は少ないので、それらの処理回数を制限する。ある種の平方分割
• sumの二重ループを減らす
  • xを決めると、max(max(U),max(V))とx+y+2のどちらが大きくなるか境界となるyが決まる。よって境界となるyをlower_boundで求め、事前にとった累積和を用いて和を計算すれば、2つめのループは不要となる。

int N,M,Q;
vector<int> E[101010];
int id[101010];
int depth[101010][4];
int hoge[101010][3];
int dia[101010];
vector<int> D[101010];
vector<ll> S[101010];
int ma;
map<pair<int,int>, double> memo;

void dfs(int cur,int pre,int d,int cid,int forest) {
	id[cur]=forest;
	depth[cur][cid]=d;
	if(d>depth[ma][cid]) ma=cur;
	FORR(e,E[cur]) if(e!=pre) dfs(e,cur,d+1,cid,forest);
}

double ask(int a,int b) {
	if(memo.count({a,b})) return memo[{a,b}];
	
	int md=max(D[a].back(),D[b].back());
	
	ll tot=0;
	FORR(r,D[a]) {
		int mi=md-(r+1);
		int x=lower_bound(ALL(D[b]),mi)-D[b].begin();
		tot += 1LL*x*md + (D[b].size()-x)*1LL*(r+1)+S[b].back()-S[b][x];
	}
	
	
	return memo[{a,b}]=tot*1.0/D[a].size()/D[b].size();
	
}

void solve() {
	int i,j,k,l,r,x,y; string s;
	
	cin>>N>>M>>Q;
	FOR(i,M) {
		cin>>x>>y;
		E[x-1].push_back(y-1);
		E[y-1].push_back(x-1);
	}
	
	x=0;
	FOR(i,N) if(id[i]==0) {
		++x;
		ma=i;
		dfs(i,-1,0,0,x);
		hoge[x][0]=ma;
		dfs(ma,-1,0,1,x);
		hoge[x][1]=ma;
		dfs(ma,-1,0,2,x);
		dia[x]=depth[hoge[x][0]][2];
	}
	
	FOR(i,N) D[id[i]].push_back(max(depth[i][1],depth[i][2]));
	FOR(i,N+1) if(D[i].size()) {
		sort(ALL(D[i]));
		S[i].push_back(0);
		FORR(r,D[i]) S[i].push_back(S[i].back()+r);
	}
	
	
	while(Q--) {
		cin>>x>>y;
		x=id[x-1];
		y=id[y-1];
		
		if(x==y) {
			_P("-1\n");
			continue;
		}
		
		if(D[x].size()>D[y].size()) swap(x,y);
		_P("%.12lf\n",ask(x,y));
	}
}

まとめ

なんとか解ききれてよかった。