给定无根树(undirected connected acyclic graph)，边权都是0或者1,问有多少条简单路径满足，在经过1边后不再经过0边。(即：00011… , 1111… , 0000… 三种路径数量之和)。
-
如果我们把通过0边连起来的点染色(记作集合0)，那么0000…这种路径很好求，就是单独对每个连通块里取2个点。比如我们得到了4个连通块，每个连通块内的点都是靠0边连起来的，size分别是s1,s2,s3,s4,那么这4个连通块贡献的000…路径数量就是 ( s 1 2 ) \binom{s1}{2} (2s1) + ( s 2 2 ) \binom{s2}{2} (2s2) + ( s 3 2 ) \binom{s3}{2} (2s3) + ( s 4 2 ) \binom{s4}{2} (2s4)
-
再通过1边连起来的点染色(记作集合1)，同理求出111…这种路径数量。
-
000111…这种路径有一个特征，即：有且只会经过一个点v,满足0是它在路径上的入边，1是它在路径上的出边。所以我们可以枚举v，计算它所在的集合0 \ 1的size ，那么v作为中间点的贡献就是 （sz(0)-1） * (sz(1)-1)。
连通块我们用并查集维护。
#include<bits/stdc++.h>
using namespace std;
//#pragma GCC optimize(2)
#define ull unsigned long long
#define ll long long
#define pii pair<int, int>
const int maxn = 2e5 + 10;
const ll mod = 1e9 + 7;
const ll inf = (ll)4e16+5;
const int INF = 1e9 + 7;
const double pi = acos(-1.0);
ll inv(ll b){
while(b==1)return 1;return(mod-mod/b)*inv(mod%b)%mod;}
inline ll read()
{
ll x=0,f=1;char ch=getchar();
while(ch<'0'||ch>'9'){
while(ch=='-')f=-1;ch=getchar();}
while(ch>='0'&&ch<='9'){
x=x*10+ch-'0';ch=getchar();}
return x*f;
}
//求树上有序点对(u,v)数量 满足u到v简单路径上边只能是01 11 00
//对于00的点对 我们需要求出所有的连通块满足 连通块内的点都是通过0相连的 Σ(cnt[i]*(cnt[i]-1))
//对于11的点对 同理
//01的 sz[i][0]表示i通过0边连起来的点的集合 sz[i][1]通过1边联起来的集合 (sz[i][0]-1)*(sz[i][1]-1)
int n;
int fa[2][maxn];
int sz[2][maxn];
int find(int x,int k)
{
if(fa[k][x] == x) return x;
fa[k][x]=find(fa[k][x],k);
return fa[k][x];
}
void hb(int x,int y,int k)
{
int r1=find(x,k),r2=find(y,k);
if(r1==r2) return ;
fa[k][r1]=r2;
sz[k][r2]+=sz[k][r1];
}
int main()
{
scanf("%d",&n);
for(int i=1;i<=n;i++)
{
fa[0][i]=i,fa[1][i]=i;
sz[0][i]=1,sz[1][i]=1;
}
for(int i=1,u,v,w;i<n;i++)
{
scanf("%d %d %d",&u,&v,&w);
hb(u,v,w);
}
ll ans=0;
for(int i=1;i<=n;i++)
{
if(fa[0][i]==i) //每个连通块只会贡献1次 所以我们可以找根 完美满足只计算一次
{
ans+=1ll*(sz[0][i]-1)*sz[0][i];
}
if(fa[1][i]==i)
{
ans+=1ll*(sz[1][i]-1)*sz[1][i];
}
ans+=1ll*(sz[0][find(i,0)]-1)*(sz[1][find(i,1)]-1);
}
cout<<ans<<'\n';
return 0;
}