【HDU4303】Hourai Jeweled

题意

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

分析

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

     树形DP。

#include <cstdio>
#include <cstring>
#include <iostream>
#include <algorithm>
#include <map>
using namespace std;
const int maxn=300000+10;
int head[maxn],Next[2*maxn],to[2*maxn],col[2*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=0,allnum=0;
    num[u]=1;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=0;
    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/2;
}
int main(){
    while(scanf("%d",&n)!=EOF){
        ans=0;
        memset(head,0,sizeof(head));
        sz=0;
        for(int i=1;i<=n;i++)
            scanf("%d",&val[i]);
        int a,b,c;
        for(int i=1;i<n;i++){
            scanf("%d%d%d",&a,&b,&c);
            add_edge(a,b,c);
            add_edge(b,a,c);
        }
        dp(1,0,0);
        printf("%lld
",ans);
    }
return 0;
}