Codeforces ECR #064 : D. 0-1-Tree

ECRのDにして結構面倒な問題。
https://codeforces.com/contest/1156/problem/D

問題

木を成すグラフが与えられる。
各辺には0,1のいずれかの値が設定されている。

頂点対(u,v)のうち、その最短パスが、最初0を設定された辺を通った後1を設定された辺を通るケースは何通りか。

解法

2つのUnion-Findを用いる。
まず、0が設定された辺で結ばれた頂点をUnion-Findで結んでいこう。
同一グループ内の頂点は、0が設定された辺だけで移動できるので、グループ内の頂点対の分を解に加える。

また、その際各頂点は重みとしてグループ内の頂点数を持つこととする。
これは、各頂点に0が設定された辺だけで移動できる点の数を示す。
次に、もう一つのUnion-Find構造で、1が設定された辺を連結する。
そして、各頂点に関し、(頂点の重み)×(グループ内の頂点数-1)の総和を解に加えよう。
これは、1の辺を1つ以上通って到達できる頂点の数を掛け合わせることに相当する。

int N;
int A[202020],B[202020],C[202020];
int W[202020],W2[202020];
ll dp[202020][2];
ll ret;

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[2];


void solve() {
	int i,j,k,l,r,x,y; string s;
	
	cin>>N;
	FOR(i,N-1) {
		cin>>A[i]>>B[i]>>C[i];
		A[i]--;
		B[i]--;
		if(C[i]==0) uf[0](A[i],B[i]);
	}
	
	ll ret=0;
	FOR(i,N) {
		W[i]=uf[0].count(i);
		if(W[i]>1 && uf[0][i]==i) ret+=1LL*W[i]*(W[i]-1);
	}
	FOR(i,N-1) if(C[i]==1) {
		uf[1](A[i],B[i]);
	}
	FOR(i,N) {
		ret+=1LL*W[i]*(uf[1].count(i)-1);
	}
	
	
	
	cout<<ret<<endl;
}

まとめ

実装は簡単だけど、実装にちょっと迷う問題。