hdu 4303 树形dp

#include<iostream>
#include<vector>
#include<algorithm>
#include<cstdio>
#include<queue>
#include<stack>
#include<string>
#include<map>
#include<set>
#include<cmath>
#include<cassert>
#include<cstring>
#include<iomanip>
using namespace std;

#ifdef _WIN32
#define i64 __int64
#define out64 "%I64d\n"
#define in64 "%I64d"
#else
#define i64 long long
#define out64 "%lld\n"
#define in64 "%lld"
#endif
/************ for topcoder by zz1215 *******************/
#define FOR(i,a,b)      for( int i = (a) ; i <= (b) ; i ++)
#define FFF(i,a)        for( int i = 0 ; i < (a) ; i ++)
#define FFD(i,a,b)      for( int i = (a) ; i >= (b) ; i --)
#define S64(a)          scanf(in64,&a)
#define SS(a)           scanf("%d",&a)
#define LL(a)           ((a)<<1)
#define RR(a)           (((a)<<1)+1)
#define pb              push_back
#define CL(Q)           while(!Q.empty())Q.pop()
#define MM(name,what)   memset(name,what,sizeof(name))
#define read            freopen("in.txt","r",stdin)
#define write           freopen("out.txt","w",stdout)

const int inf = 0x3f3f3f3f;
const i64 inf64 = 0x3f3f3f3f3f3f3f3fLL;
const double oo = 10e9;
const double eps = 10e-9;
const double pi = acos(-1.0);
const int maxn = 300011;

struct zz
{
	int from;
	int to;
	int c;
}zx;

vector<zz>g[maxn];
i64 s[maxn];
int tot[maxn];
int ex[maxn];
int up[maxn];
int f[maxn];
int col[maxn];
bool vis[maxn];
map<int,int>mp[maxn];
i64 stot[maxn];
i64 sex[maxn];
i64 sup[maxn];
map<int,i64>smp[maxn];
int n;

void dfs(int now)
{
	int to;
	for(int i=0;i<g[now].size();i++)
	{
		to = g[now][i].to;
		if(!vis[to])
		{
			vis[to]=true;
			f[to]=now;
			col[to]=g[now][i].c;
			dfs(to);
		}
	}
	return ;
}

void df2(int now)
{
	int to;
	for(int i=0;i<g[now].size();i++)
	{
		to = g[now][i].to;
		if(to!=f[now])
		{
			df2(to);
		}
	}
	for(int i=0;i<g[now].size();i++)
	{
		to = g[now][i].to;
		if(to!=f[now])
		{
			tot[now]+=ex[to]+1;
			mp[now][col[to]]+=ex[to]+1;
			stot[now]+=sex[to]+ex[to]*s[now]+s[to]+s[now];
			smp[now][col[to]]+=sex[to]+ex[to]*s[now]+s[to]+s[now];
			if(col[to]!=col[now])
			{
				ex[now]+=ex[to]+1;
				sex[now]+=sex[to]+ex[to]*s[now]+s[to]+s[now];
			}
		}
	}
	return ;
}

void df3(int now)
{
	int to;
	for(int i=0;i<g[now].size();i++)
	{
		to = g[now][i].to;
		if(to!=f[now])
		{
			if(col[to]!=col[now])
			{
				up[to]+=up[now];
				sup[to]+=sup[now]+up[now]*s[to];
			}
			up[to]+=tot[now]-mp[now][col[to]]+1;
			sup[to]+=stot[now]-smp[now][col[to]]+s[to]*(tot[now]-mp[now][col[to]])+s[now]+s[to];
			df3(to);
		}
	}
	return ;
}

void start()
{
	for(int i=1;i<=n;i++)
	{
		tot[i]=ex[i]=up[i]=0;
		stot[i]=sex[i]=sup[i]=0;
	}
	for(int i=1;i<=n;i++)
	{
		vis[i]=false;
	}
	f[1]=col[1]=0;
	vis[1]=true;
	dfs(1);
	df2(1);
	df3(1);
	return ;
}

int main()
{
	while(cin>>n)
	{
		int now,to,c;
		for(int i=1;i<=n;i++)
		{
			g[i].clear();
			mp[i].clear();
			smp[i].clear();
			S64(s[i]);
		}
		for(int i=1;i<=n-1;i++)
		{
			SS(now);
			SS(to);
			SS(c);
			zx.from = now;
			zx.to = to;
			zx.c = c;
			g[now].pb(zx);
			swap(zx.from,zx.to);
			g[to].pb(zx);
		}
		start();
		i64 ans=0;
		for(int i=1;i<=n;i++)
		{
			ans+=stot[i]+sup[i];
		}
		ans/=2;
		printf("%I64d\n",ans);
	}
    return 0;
}