hdu 4303 Hourai Jeweled(树形dp)

题意：给出一棵树，树上每个节点有权值，树上每条边有一种颜色。如果一条路径中没有两条相邻的边的颜色相同，那么就称这条路径为gorgeous roads，求所有gorgeous roads上的点的权值和。

思路：这题写的真的很蛋疼啊，思路其实还算好想。我们用dp[u]表示可以从u的子节点到u的权值和，用cntp[u]表示可以从u的子节点到u的点的个数，对于一个结点u来说，如果有能到达u的路径，那么结果就要加上dp[v]+val[u]*cntp[v](val[u]为u的权值)，另外还有这个点和其他已经访问过的点的路径的组合，因此还要再加上cntp[v]*已经访问过的权值和(除去边颜色相同的情况)+dp[v]*已经访问过的点的数量(除去边颜色相同的情况)+val[u]*(路径中经过u的次数)。对于颜色的处理方法，先用pa[u]记录u的父边的颜色，如果它的子节点的边的颜色和pa[u]相同，那么dp[u]就不加上它，另外再用两个数组存u的每种颜色的边连的节点的子树的权值和还有每种颜色的边连的节点的子树的节点数就行了。

 

代码：

 

#include<iostream>
#include<cstdio>
#include<cstring>
#include<string>
#include<algorithm>
#include<map>
#include<queue>
#include<set>
#include<stack>
#include<cmath>
#include<vector>
#define inf 0x3f3f3f3f
#define Inf 0x3FFFFFFFFFFFFFFFLL
#define eps 1e-9
#define pi acos(-1.0)
using namespace std;
typedef long long ll;
const int maxn=300000+10;
const int maxv=100000+10;
struct Edge
{
    int v,w,next;
};
Edge edges[maxn<<1];
int head[maxn],nEdge,n;
int val[maxn],pa[maxn],cntp[maxn],S[maxn],c;
ll dp[maxn],cw[maxv],cv[maxv],ans;
void AddEdge(int u,int v,int w)
{
    nEdge++;
    edges[nEdge].v=v;
    edges[nEdge].w=w;
    edges[nEdge].next=head[u];
    head[u]=nEdge;
}
void dfs(int u,int fa)
{
    dp[u]=val[u];cntp[u]=1;
    ll sum=0;
    for(int k=head[u];k!=-1;k=edges[k].next)
    {
        int v=edges[k].v;
        if(v==fa) continue;
        pa[v]=edges[k].w;
        dfs(v,u);
    }
    c=0;
    for(int k=head[u];k!=-1;k=edges[k].next)
    {
        int v=edges[k].v;
        if(v==fa) continue;
        if(!cv[edges[k].w]) S[c++]=edges[k].w;
        cntp[u]+=cntp[v];
        cv[edges[k].w]+=cntp[v];
        cw[edges[k].w]+=dp[v];
        if(edges[k].w!=pa[u])
           dp[u]+=dp[v]+(ll)val[u]*cntp[v];
        ans+=dp[v]+(ll)val[u]*cntp[v];
        sum+=dp[v];
        ans+=(sum-cw[edges[k].w])*cntp[v]+dp[v]*(cntp[u]-1-cv[edges[k].w])+(ll)val[u]*cntp[v]*(cntp[u]-1-cv[edges[k].w]);
    }
    cntp[u]-=cv[pa[u]];
    while(c>0)
    {
        cv[S[--c]]=0;
        cw[S[c]]=0;
    }
}
int main()
{
    //freopen("in.txt","r",stdin);
    //freopen("out.txt","w",stdout);
    while(~scanf("%d",&n))
    {
        memset(head,0xff,sizeof(head));
        memset(cw,0,sizeof(cw));
        memset(cv,0,sizeof(cv));
        nEdge=-1;
        for(int i=1;i<=n;++i)
          scanf("%d",&val[i]);
        int u,v,w;
        for(int i=1;i<n;++i)
        {
            scanf("%d%d%d",&u,&v,&w);
            AddEdge(u,v,w);
            AddEdge(v,u,w);
        }
        ans=0;pa[1]=0;
        dfs(1,-1);
        printf("%I64d\n",ans);
    }
    return 0;
}