BZOJ3697【点分治】

/* I will wait for you */

#include<cstdio>
#include<cstdlib>
#include<cstring>
#include<cmath>
#include<ctime>
#include<algorithm>
#include<iostream>
#include<fstream>
#include<vector>
#include<queue>
#include<deque>
#include<set>
#include<map>
#include<string>

typedef long long LL;
typedef unsigned long long ULL;

using namespace std;

const int maxn=200010;
const int maxm=1010;
const int maxs=26;
const int INF=1<<29;
const int P=1000000007;
const double error=1e-9;

struct edge{ int v,w,next; }e[maxn];

int head[maxn],t[maxn],size[maxn],mx[maxn],del[maxn],dis[maxn],deep[maxn];
int n,sum,root,mxdeep,cnt=0;
LL f[maxn][2],g[maxn][2],ans;

void add(int u,int v,int w)
{
	e[cnt]=(edge){v,w,head[u]};head[u]=cnt++;
}

void findroot(int u,int p)
{
	size[u]=1;mx[u]=0;
	for(int i=head[u];i!=-1;i=e[i].next)
	{
		int v=e[i].v;
		if(!del[v]&&v!=p)
		{
			findroot(v,u);
			size[u]+=size[v];
			mx[u]=max(mx[u],size[v]);
		}
	}
	mx[u]=max(mx[u],sum-size[u]);
	if(mx[u]<mx[root]) root=u;
}

void dfs(int u,int p)
{
	g[dis[u]][t[dis[u]]>0]++;
	mxdeep=max(mxdeep,deep[u]);
	t[dis[u]]++;
	for(int i=head[u];i!=-1;i=e[i].next)
	{
		int v=e[i].v;
		if(!del[v]&&v!=p)
		{
			dis[v]=dis[u]+e[i].w;
			deep[v]=deep[u]+1;dfs(v,u);
		}
	}
	t[dis[u]]--;
}
	
void solve(int u)
{
	del[u]=1;f[n][0]=1;
	int mxsum=0;
	
	for(int i=head[u];i!=-1;i=e[i].next)
	{
		int v=e[i].v;
		if(!del[v])
		{
			deep[v]=1;dis[v]=n+e[i].w;
			mxdeep=0;dfs(v,u);mxsum=max(mxsum,mxdeep);
			ans+=(f[n][0]-1)*g[n][0];
			for(int i=-mxdeep;i<=mxdeep;i++)
				ans+=f[n+i][0]*g[n-i][1]+f[n+i][1]*g[n-i][0]+f[n+i][1]*g[n-i][1];
			for(int i=-mxdeep;i<=mxdeep;i++)
				f[n+i][0]+=g[n+i][0],f[n+i][1]+=g[n+i][1],g[n+i][0]=g[n+i][1]=0;
		}
	}
	
	for(int i=-mxsum;i<=mxsum;i++) f[n+i][0]=f[n+i][1]=0;
	
	for(int i=head[u];i!=-1;i=e[i].next)
	{
		int v=e[i].v;
		if(!del[v])
		{
			root=0;sum=size[v];
			findroot(v,0);solve(root);
		}
	}
}	

int main()
{
	scanf("%d",&n);
	memset(head,-1,sizeof(head));
	for(int i=1;i<n;i++)
	{
		int u,v,w;scanf("%d%d%d",&u,&v,&w);
		w=w?1:-1;add(u,v,w);add(v,u,w);
	}
	
	sum=mx[0]=n;
	findroot(1,0);solve(root);
	
	printf("%lld\n",ans);
	return 0;
}