Hourai Jeweled(HDU4303)

题意

辉夜有一支呈树型的蓬莱玉枝，上面有n个点，有n-1条边连接它们。不同的边颜色不同。从x到y的一条路，如果满足相邻边的颜色都不相同，那么这就是一条adjacent的路，这条路的adjacent值是这条路上所有节点的权值总和。
求这支蓬莱玉枝的总adjacent值。

分析

对于所有点计算贡献的方法吧。
考虑这个点到点颜色各不相同的边的个数，同时保证不在一颗子树内(也就是保证连出去的边没有交叉的部分)。
然后统计这个点被算进去的方案数量，乘上点权加入总答案。

觉得下方代码并不是很优雅。应该有更漂亮的写法的。

code

#include<bits/stdc++.h>
#define ll long long
#define pii pair<int,int>
#define mp make_pair
#define A first
#define B second
#define M 300005
#define N 100005
using namespace std;

void read(ll &x){
    x=0; char c=getchar();
    for (; c<'0'; c=getchar());
    for (; c>='0'; c=getchar())x=(x<<3)+(x<<1)+(c^'0'); 
}
ll f[M],h[M],sum[N],cnt[N],res,v[M];
int nx[M];
int ecnt;
struct ed{
    int A,B,nx; 
}e[M<<1];
void add(int x,int y,int v){
    e[ecnt]=(ed){y,v,nx[x]};
    nx[x]=ecnt++;
    e[ecnt]=(ed){x,v,nx[y]};
    nx[y]=ecnt++;
}
void dfs(int fr,int x,int co){
    int i; 
    f[x]=v[x]; h[x]=1;
    for (i=nx[x]; i; i=e[i].nx)if (e[i].A!=fr){
        int to=e[i].A;
        dfs(x,to,e[i].B);
        if (e[i].B!=co)f[x]+=f[to]+v[x]*h[to],h[x]+=h[to];
    }
    ll Sum=0;
    ll Cnt=0;
    ll Res=0;
    for (i=nx[x]; i; i=e[i].nx)if (e[i].A!=fr){
        int to=e[i].A,col=e[i].B;
        Res+=f[to]*(Cnt-cnt[col])+h[to]*(Sum-sum[col]);
        Res+=v[x]*h[to]*(Cnt-cnt[col]);
        Res+=f[to]+v[x]*h[to];
        cnt[col]+=h[to]; sum[col]+=f[to];
        Cnt+=h[to]; Sum+=f[to];
    }
    for (i=nx[x]; i; i=e[i].nx)sum[e[i].B]=cnt[e[i].B]=0;
    res+=Res;
}

int main(){
    int i,n;
    ll x,y,z;
    for (; scanf("%d",&n)!=EOF; ){
        memset(nx,0,sizeof(nx));
        for (i=1; i<=n; i++)read(v[i]);
        ecnt=1;
        for (i=1; i<n; i++){
            read(x); read(y); read(z);add(x,y,z);
        }
        res=0; dfs(0,1,0);
        printf("%lld\n",res);
    }
    return 0;   
}