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 ;

 }

巴特西

HDU 4303 Hourai Jeweled（树形DP）

最新文章

热门文章