Codeforces #212 Div2. D. Fools and Foolproof Roads

アプローチはあっていたのに凡ミスした…。
何気にひっかけポイントが多い問題。
http://codeforces.com/contest/362/problem/D

問題

N個の点・M本の無向辺からなるグラフ構造が与えられる。
各辺には長さL[i]というパラメータがある。
辺を通じて互いに行き来可能な点をregionと呼ぶ。

ここで、以下の手順でP個の辺を追加したい。

• 2つの異なる点の間に辺を加える。同じ点の対に対し複数本辺があってもよい。
• 加える辺の長さは、2つの点が同じregionに含まれるなら1000、異なるregionに含まれるなら、両regionに含まれる辺の長さの総和+1か、10^9の小さい方である。

最終的に、regionの個数をQ個にしたい。また、追加する辺の長さを最小化したい。
regionをQ個にできるか、またQ個にできるときに辺を追加する手順を答えよ。

解法

まずunion-findでregionを判定する。また、region内の辺の長さの総和を求める。
ここで、最終的にregionをQ個にできるか先に判定する。
以下のケースではQ個にできない。

• もともとregionがQ個未満しかない
• regionにP個辺をくっつけてP個regionを減らしてもQ個より大きい
• 初期状態でN点がすべてのregionで、かつP==0・Q==Nであり、1個でも辺を張るとregionがN-1個になってしまう

それ以外のケースではQ個にできる。
後は以下の手順で点をつなげていけばよい。

• region数がQ個より多いなら、辺の総長の短い2つのregionをくっつける
• region数がQ個なら、同じregion内の2点をくっつける

int N,M,P,Q;
int X[100001],Y[100001];
ll L[100001];

const int ufmax=1000001;
int ufpar[ufmax],ufrank[ufmax];
void UF_init(){int i; FOR(i,ufmax) { ufpar[i]=i; ufrank[i]=0; } }
int UF_find(int x) {	return (ufpar[x]==x)?(x):(ufpar[x] = UF_find(ufpar[x]));}
void UF_unite(int x,int y) {
	x = UF_find(x); y = UF_find(y);
	if(x==y) return;
	if(ufrank[x]<ufrank[y]) ufpar[x]=y;
	else {ufpar[y]=x; if(ufrank[x]==ufrank[y]) ufrank[x]++;}
}

map<int,ll> MM;
set<pair<ll,int> > S;
int rep=-1;

void solve() {
	int f,i,j,k,l,x,y;
	
	cin>>N>>M>>P>>Q;
	UF_init();
	FOR(i,M) {
		cin>>X[i]>>Y[i]>>L[i];
		UF_unite(X[i],Y[i]);
	}
	FOR(i,N) MM[UF_find(i+1)]=0;
	FOR(i,N) if(i+1!=UF_find(i+1)) rep=i+1;
	FOR(i,M) MM[UF_find(X[i])]+=L[i];
	
	if((int)MM.size()-P>Q) return _P("NO\n");
	if(MM.size()<Q) return _P("NO\n");
	if(MM.size()==N && P>0 && N==Q) return _P("NO\n");
	ITR(it,MM) S.insert(make_pair(it->second,it->first));
	
	_P("YES\n");
	FOR(i,P) {
		if(S.size()==Q) {
			_P("%d %d\n",rep,UF_find(rep));
		}
		else {
			set<pair<ll,int> >::iterator itr=S.begin();
			pair<ll,int> p1=*itr++;
			pair<ll,int> p2=*itr;
			S.erase(p1);
			S.erase(p2);
			S.insert(make_pair(p1.first+p2.first+min(1000000000LL,p1.first+p2.first+1),p1.second));
			_P("%d %d\n",p1.second,p2.second);
			UF_unite(p1.second,p2.second);
			if(p1.second!=UF_find(p1.second)) rep=p1.second;
			if(p2.second!=UF_find(p2.second)) rep=p2.second;
		}
		
	}
}

まとめ

なんかハフマン符号の処理みたいだ。
ソート済みの数列を作るのにsetを使ってるけど、他にいい方法あるのかな？
priority_queueも使える？