Codeforces #363 Div1 B. Fix a Tree

コーナーケースがありそうでヒヤッとしたけど大丈夫だった。
http://codeforces.com/contest/698/problem/B

問題

N頂点N辺のグラフが与えられる。
頂点iは頂点P[i]と辺を結ぶ。なおi==P[i]でも良い。

与えられたグラフが連結根付木となるグラフとなるようにしたい。（根はi==P[i]となる頂点i）
入力のP[i]が連結根付木を成していない場合、連結根付木になるよういくつかの要素を書き換えたい。
最小何要素を書き換えればよいか。

解法

まず根を決めよう。
i==P[i]となる頂点が１個でもあれば、それを根とする。
１個もなければ、まずは根を保留する。

次に各iについて、Union-Findの要領でiとP[i]を接続していく。
なお、グラフが木を成すので、根以外の頂点ではiとP[i]を結ぶことで非連結成分が連結されなければならない。
逆にiとP[i]を接続する前に既に連結していた場合、このP[i]は書き換え対象となる。
根が決まっていたら、P[i]をその根に変更しよう。
まだ根が決まっていなかったらP[i]=iとし、その頂点を根とすることにしよう。

あとはこの処理を全頂点に繰り返す。

template<int um> class UF {
	public:
	vector<int> par,rank;
	UF() {rank=vector<int>(um,0); for(int i=0;i<um;i++) par.push_back(i);}
	int operator[](int x) {return (par[x]==x)?(x):(par[x] = operator[](par[x]));}
	int operator()(int x,int y) {
		if((x=operator[](x))==(y=operator[](y))) return x;
		if(rank[x]>rank[y]) return par[x]=y;
		rank[x]+=rank[x]==rank[y]; return par[y]=x;
	}
};
UF<500000> uf;

int N,fix;
int A[202020];
int isroot[202020];
int rootcand=-1;

void solve() {
	int i,j,k,l,r,x,y; string s;
	
	cin>>N;
	
	FOR(i,N) {
		cin>>A[i];
		A[i]--;
		if(i==A[i]) rootcand=i;
	}
	FOR(i,N) {
		
		if(i==A[i]) {
			if(i!=rootcand) {
				fix++;
				A[i]=rootcand;
			}
		}
		else if(uf[i]==uf[A[i]]) {
			if(rootcand==-1) rootcand=i;
			fix++;
			A[i]=rootcand;
		}
		uf(i,A[i]);
	}
	_P("%d\n",fix);
	FOR(i,N) _P("%d%c",A[i]+1,(i==N-1)?'\n':' ');
	
}

まとめ

こういうconstructionに近い問題は何か見落としてそうで怖い。