【HDU4303】Hourai Jeweled

题意

有一棵n个结点的树，每个结点都有一个值，没一条边都有一个颜色。如果某条路径上，相邻的边颜色不同，那么把这路径上所有的点的值加起来。 输出所有符合条件的路径上值的和。 n<=300000。

分析

场上读错题了。当时以为是路径上所有的边都不能有相同的颜色，雾。。。

树形DP。

 #include <cstdio>
 #include <cstring>
 #include <iostream>
 #include <algorithm>
 #include <map>
 using namespace std;
 const int maxn=+;
 int head[maxn],Next[*maxn],to[*maxn],col[*maxn],val[maxn];
 int n,sz;
 long long ans;
 void add_edge(int a,int b,int c){
     sz++;
     to[sz]=b;
     col[sz]=c;
     Next[sz]=head[a];
     head[a]=sz;
 }
 long long f[maxn],num[maxn];
 void dp(int u,int fa,int co){
     long long alltmp=,allnum=;
     num[u]=;f[u]=val[u];
     map<int,long long>Mdp;
     map<int,long long>Mnum;

     for(int i=head[u];i;i=Next[i]){
         int v=to[i];
         if(v!=fa){
             dp(v,u,col[i]);
             if(co!=col[i]){
                 f[u]+=f[v]+val[u]*num[v];
                 num[u]+=num[v];
             }
             ans+=f[v]+val[u]*num[v];
             alltmp+=f[v]+val[u]*num[v];
             allnum+=num[v];
             Mnum[col[i]]+=num[v];
             Mdp[col[i]]+=f[v]+val[u]*num[v];
         }
     }
     long long temp=;
     for(int i=head[u];i;i=Next[i]){
         if(to[i]!=fa){
             temp+=f[to[i]]*(allnum-Mnum[col[i]])+num[to[i]]*(alltmp-Mdp[col[i]]);
         }
     }
     ans+=temp/;
 }
 int main(){
     while(scanf("%d",&n)!=EOF){
         ans=;
         memset(head,,sizeof(head));
         sz=;
         for(int i=;i<=n;i++)
             scanf("%d",&val[i]);
         int a,b,c;
         for(int i=;i<n;i++){
             scanf("%d%d%d",&a,&b,&c);
             add_edge(a,b,c);
             add_edge(b,a,c);
         }
         dp(,,);
         printf("%lld\n",ans);
     }
 return ;
 }