D. 0-1-Tree(染色+计数)

题目

题意：

    给定一棵带权树，权值为0或1。要求求出满足一下条件的点对(x,y)的数量。条件为从x到y的路径中不可以存在经过了权值为1的边然后再经过权值为0的边。
     $2≤n≤200000,1≤x_i,y_i≤n , 0≤c_i≤1$

分析：

    经过分析，可以得到满足条件的点对，要么路径上为全1，要么全0，要么先0再1。这样的话我们就可以去求只有边权为1的连通块和边权为0的连通块，即染色。染色之后，全0和全1就可以轻松算出。对于先0再1的，我们考虑一些点，它既在边权为0的连通块中，又在边权为1的连通块中，那么它所连接的两个01连通块就可以通过这一点形成答案，即先0后1的情况，注意计算时排除这个点。

#include <iostream>
#include <vector> 
using namespace std;

typedef long long ll;

struct node{
	int num,val;
	node(int a,int b)
	{
		num = a;
		val = b;
	}
};

vector<node> g[200005];
int color1[200005],color2[200005],size1[200005],size2[200005];
int c1 = 1,c2 = 1;

void dfs1(int x,int fa)
{
	if( color1[fa] == 0 && fa != 0 )
	{
		color1[fa] = c1;
		size1[c1] ++;
	}
	for (int i = 0; i < g[x].size(); i++)
	{
		node t = g[x][i];
		if( color1[t.num] ) continue;
		if( t.val == 1 )
		{
			dfs1(t.num,x);
		}
	}
	if( color1[x] == 0 && fa != 0 )
	{
		color1[x] = c1;
		size1[c1] ++;
	}
}

void dfs2(int x,int fa)
{
	if( color2[fa] == 0 && fa != 0 )
	{
		color2[fa] = c2;
		size2[c2] ++;
	}
	for (int i = 0; i < g[x].size(); i++)
	{
		node t = g[x][i];
		if( color2[t.num] ) continue;
		if( t.val == 0 )
		{
			dfs2(t.num,x);
		}
	}
	if( color2[x] == 0 && fa != 0 )
	{
		color2[x] = c2;
		size2[c2] ++;
	}
}

int main()
{
	ios::sync_with_stdio(false);
	cin.tie(0);
	int n;
	cin >> n;
	for (int i = 1; i < n; i++)
	{
		int x,y,v;
		cin >> x >> y >> v;
		g[x].push_back(node(y,v));
		g[y].push_back(node(x,v)); 
	}
	for (int i = 1; i <= n; i++)
	{
		if( color1[i] == 0 ) 
		{
			dfs1(i,0);
			if( size1[c1] > 0 ) c1 ++;
		}
		if( color2[i] == 0 ) 
		{
			dfs2(i,0);
			if( size2[c2] > 0 ) c2 ++;
		}
	}
	ll ans = 0;
	//cout << c1 << ' ' << size1[c1-1] << ' ' << c2 << ' ' << size2[c2-1] << '\n';
	for (int i = 1; i < c1; i++)
	{
		ans += (ll)( size1[i] - 1 ) * size1[i];
	}
	for (int i = 1; i < c2; i++)
	{
		ans += (ll)( size2[i] - 1 ) * size2[i];
	}
	for (int i = 1; i <= n; i++)
	{
		if( color1[i] != 0 && color2[i] != 0 )
		{
			ans += (ll)(size2[color2[i]] - 1) * ( size1[color1[i]] - 1 ); 
		}
	}
	cout << ans << '\n';
	return 0;
}