yukicoder : No.1293 2種類の道路

★3.5以下は何とか解き切った。
https://yukicoder.me/problems/no/1293

問題

N頂点M辺のグラフが与えられる。
辺はすべて無向辺で、2種類のタイプがある。

1種類目の辺を使って移動後、2種類目の辺で移動した場合、移動可能な頂点対の組み合わせを求めよ。

解法

両者の辺で個別にUnion-Findで連結関係を求めよう。
setなど使い、片方のUnion-FIndで連結される要素に対し、対応するもう片方のUnion-Findの和集合の要素数を掛け合わせていこう。

int N,D,W;

template<int um> class UF {
	public:
	vector<int> par,rank,cnt;
	UF() {par=rank=vector<int>(um,0); cnt=vector<int>(um,1); for(int i=0;i<um;i++) par[i]=i;}
	void reinit() {int i; FOR(i,um) rank[i]=0,cnt[i]=1,par[i]=i;}
	int operator[](int x) {return (par[x]==x)?(x):(par[x] = operator[](par[x]));}
	int count(int x) { return cnt[operator[](x)];}
	int operator()(int x,int y) {
		if((x=operator[](x))==(y=operator[](y))) return x;
		cnt[y]=cnt[x]=cnt[x]+cnt[y];
		if(rank[x]>rank[y]) return par[x]=y;
		rank[x]+=rank[x]==rank[y]; return par[y]=x;
	}
};
UF<202020> uf,uf2;
set<int> C[202020];

void solve() {
	int i,j,k,l,r,x,y; string s;
	
	cin>>N>>D>>W;
	FOR(i,D) {
		cin>>x>>y;
		uf(x-1,y-1);
	}
	FOR(i,W) {
		cin>>x>>y;
		uf2(x-1,y-1);
	}
	
	FOR(i,N) {
		C[uf2[i]].insert(uf[i]);
	}
	
	ll ret=0;
	FOR(i,N) if(uf2[i]==i) {
		ll a=uf2.count(i);
		ll b=0;
		FORR(e,C[i]) b+=uf.count(e);
		ret+=a*b;
	}
	ret-=N;
	cout<<ret<<endl;
	
}

まとめ

本番無駄にマージテクを使ってしまった。