Artifact 8b2db66f69d32595c830db74f19b72c50763313f

• File SRM/424/2C.cpp
  • 2011-02-23 09:21:16 - part of checkin [4fd800b3a8] on branch trunk - Copied from private svn repository. (user: kinaba) [annotate]

#include <vector>
#include <string>
#include <set>
using namespace std;

struct BestRoads
{
	vector<int> numberOfRoads(vector<string> roads, int M)
	{
		int N = roads.size();
		vector<int> ans(N);

		// Kruskal
		typedef pair<int,int> road;
		set<road> Q;
		for(int i=0; i<N; ++i)
			for(int j=i+1; j<N; ++j)
				if( roads[i][j]=='Y' )
					Q.insert( road(i,j) );

		vector<int> p(N,-1), s(N, 1); // Union-Find
		int div=N, red=M-N+1;

		for(set<road>::iterator it=Q.begin(); it!=Q.end(); ++it)
		{
			road r = *it;
			int a=r.first;  while(p[a]!=-1) a=p[a];
			int b=r.second; while(p[b]!=-1) b=p[b];
			if( a != b )
			{
				ans[r.first]++, ans[r.second]++, div--;
				if(s[a]<s[b]) p[a]=b, s[b]+=s[a];
				else          p[b]=a, s[a]+=s[b];
			}
			else if( red )
				ans[r.first]++, ans[r.second]++, red--;
		}

		return (red==0 && div==1) ? ans : vector<int>();
	}
};