HDU 4303 Hourai Jeweled（树形DP）

http://acm.hdu.edu.cn/showproblem.php?pid=4303

题意：
给出一棵树，树上的每一个节点都有一个权值，每条边有一个颜色，如果一条路径上相邻边的颜色都是不同的，那么它就是符合要求的。求出所有符合要求的路径上的节点的权值和。

思路：
num[u]表示u节点下有几条符合要求的子树路径，sum[u]表示u为起点（或者终点也可以）往子树方向符合要求的路径权值和。

如图，u的父节点颜色为1，u->v的边颜色为2，那么此时u可以和v相连，num[v]就是v保留的路径数，这些保留下来的都是和u->v颜色不同的，那么在原有的sum[v]上，因为有num[v]条路径，那么u节点可以连num[v]次，所以此时sum[u]+=sum[v]+num[v]*val[u]， num[u]+=num[v]。

但是这样的话处理了u作为端点的情况，还存在一种情况就是u的两个子节点相连。这里可以用map来记录不同颜色的路径数和不同颜色的权值总和。统一处理即可。

#include<iostream>
#include<cstdio>
#include<cstring>
#include<map>
using namespace std;
const int maxn = 300000 + 5;
typedef long long ll;

int n,tot;
int head[maxn], val[maxn],num[maxn];
ll sum[maxn];
ll ans;

struct node
{
    int v,col,next;
}e[2*maxn];


void addEdge(int u,int v,int col)
{
    e[tot].v = v;
    e[tot].col = col;
    e[tot].next = head[u];
    head[u] = tot++;
}


void dfs(int u,int fa,int prec)
{
    num[u] = 1;
    sum[u] = val[u];
    map<int,ll> mpNum, mpSum;
    ll temp = 0, tempNum = 0, tempSum = 0;
    for(int i=head[u];i!=-1;i=e[i].next)
    {
        int v = e[i].v;
        int col = e[i].col;
        if(v == fa) continue;
        dfs(v,u,col);
        temp = sum[v] + (ll)num[v]*val[u];
        if(col != prec)
        {
            num[u] += num[v];
            sum[u] += temp;
        }
        tempNum += num[v];
        tempSum += temp;
        mpNum[col] += num[v];
        mpSum[col] += temp;
    }
    ans += tempSum;
    temp = 0;
    for(int i=head[u];i!=-1;i=e[i].next)
    {
        int v = e[i].v;
        int col = e[i].col;
        if(v == fa) continue;
        temp += sum[v]*(tempNum - mpNum[col]) + num[v]*(tempSum - mpSum[col]);
    }
    ans += temp/2;
}

int main()
{
    while(~scanf("%d",&n))
    {
        ans = 0;
        tot = 0;
        memset(head,-1,sizeof(head));
        for(int i=1;i<=n;i++) scanf("%d",&val[i]);    
        for(int i=1;i<n;i++)
        {
            int u,v,w;
            scanf("%d%d%d",&u,&v,&w);
            addEdge(u,v,w);
            addEdge(v,u,w);
        }
        dfs(1,0,0);
        printf("%lld
",ans);
    }
    return 0;
}