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 =  + ;
 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[*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] = ;
     sum[u] = val[u];
     map<int,ll> mpNum, mpSum;
     ll temp = , tempNum = , tempSum = ;
     for(int i=head[u];i!=-;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 = ;
     for(int i=head[u];i!=-;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/;
 }

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