C++11(clang++ 3.9) 解法, 执行用时: 744ms, 内存消耗: 17912K, 提交时间: 2020-07-07 23:08:01

#include<bits/stdc++.h>
using namespace std;
const int maxn=100009;

int n;

int head[maxn],cntedge;
int nex[maxn<<1],to[maxn<<1],dist[maxn<<1];
void Addedge(int x,int y,int z){
	nex[++cntedge]=head[x];
	to[cntedge]=y;
	dist[cntedge]=z;
	head[x]=cntedge;
}

int vis[maxn],siz[maxn],mxsiz[maxn];
int getroot(int x,int fa,int tot){
	siz[x]=1;mxsiz[x]=0;
	int cen=0;
	for(int i=head[x];i;i=nex[i]){
		int v=to[i];
		if(vis[v]||v==fa)continue;
		int vcen=getroot(v,x,tot);
		if(!cen||mxsiz[vcen]<mxsiz[cen])cen=vcen;
		siz[x]+=siz[v];
		mxsiz[x]=max(mxsiz[x],siz[v]);
	}
	mxsiz[x]=max(mxsiz[x],tot-siz[x]);
	if(!cen||mxsiz[x]<mxsiz[cen])cen=x;
	return cen;
}


long long ans;
map<int,int>tong;
int q0[maxn],n0;
int q1[maxn],n1;
void dp(int x,int fa,int del){
	if(tong[del]!=0)q1[++n1]=del;
	else q0[++n0]=del;
	tong[del]++;
	for(int i=head[x];i;i=nex[i]){
		int v=to[i];
		if(vis[v]||v==fa)continue;
		if(dist[i])dp(v,x,del+1);
		else dp(v,x,del-1);
	}
	tong[del]--;
}

void calc(int x){
	map<int,int>S0,S1;
	for(int i=head[x];i;i=nex[i]){
		int v=to[i];
		if(vis[v])continue;
		n0=0;n1=0;tong.clear();
		if(dist[i])dp(v,x,1);
		else dp(v,x,-1);
		for(int j=1;j<=n0;++j){
			ans+=S1[-q0[j]];
			if(q0[j]==0)ans+=S0[0];
		}
		for(int j=1;j<=n1;++j){
			ans+=S0[-q1[j]]+S1[-q1[j]];
		}
		for(int j=1;j<=n0;++j)S0[q0[j]]++;
		for(int j=1;j<=n1;++j)S1[q1[j]]++;
	}
	ans+=S1[0];
}
void solve(int x,int tot){
	vis[x]=1;calc(x);
	for(int i=head[x];i;i=nex[i]){
		int v=to[i];
		if(vis[v])continue;
		int vsiz=(siz[v]<siz[x])?siz[v]:tot-siz[x];
		int vcen=getroot(v,x,vsiz);
		solve(vcen,vsiz);
	}
}

int main(){
	scanf("%d",&n);
	for(int i=1;i<=n-1;++i){
		int x,y,z;
		scanf("%d%d%d",&x,&y,&z);
		Addedge(x,y,z);
		Addedge(y,x,z);
	}
	
	int cen=getroot(1,0,n);
	solve(cen,n);
	
	printf("%lld\n",ans);
	return 0;
}