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;}