Artifact 8b2db66f69d32595c830db74f19b72c50763313f
#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>();
}
};