hdu 4303 Hourai Jeweled(树型DP+统计答案)

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

题意：给你一颗树，每个节点有权值，每条边有颜色，问所有满足相邻边颜色不同的路径的权值之和

分析：一开始看成求权值最大的路径，所以就是水题了，不过加个统计就得略想一下子，f[ i ]表示i的所有子节点的连到i 路径和，这里有个限制就是i到子节点的颜色与i到父节点的颜色不同，g[ i ]记录相应的边的数量，于是有f[ i ]=sum{ f[ j ]+g[ j ]*v[i] },  g[ i ]=sum{ g[ j ] }（ i和j满足颜色限制）这样，对于每个节点i有3种路径，一种以它为端点的路径，这个答案ans1=sum{ f[ j ]+g[ j ]*v[i] }这里没有颜色限制，第二种就是通过i到i的父节点，这个可以转化成父节点的问题，所以不用考虑，回溯的时候自然解决了，第三种最复杂，就是两个来自子节点的边在i点和成一条路径，这里因为需要颜色限制所以比较烦，不过，还是有简单的方法的，sumf = sum{ f[ j ]+g[ j ]*v[ i ] }, sumg=sum{ g[j ] }这里i和j没有颜色限制，对于来自每个子节点j的边我们这样统计，ans2=(sumf - sum{ f[ k ]+g[ k ]*v[i] })*g[ j ] + (sumg-sum{ g[k ] }) *f[ j ]   这里k表示子节点k到i的边与j到i的边颜色相同，至于为什么这样算自己想想就明白了，很简单，所以答案就是每个节点的ans1+ans2的和

注意使用scanf读入，cin超时啊T_T

代码：

#include <cstdio>
#include <iostream>
#include <map>
using namespace std;
#define REP(i,n) for(i=0;i<(n);++i)
#define N 666666
typedef long long ll;
int head[N],ver[N],col[N],next[N],vis[N];
ll f[N],g[N],a[N];
int edge;
ll ans;
void dp(int u,int cc)
{
    f[u]=a[u];
    g[u]=1;
    vis[u]=1;
    ll tmpf=0,tmpg=0,tmp=0;
    map<int,ll> mg,mf;
    for(int i=head[u],v;i>=0;i=next[i])
        if(!vis[v=ver[i]])
        {
            dp(v,col[i]);
            tmp=f[v]+a[u]*g[v];
            if(col[i]!=cc)
            {
                f[u]+=tmp;
                g[u]+=g[v];
            }
            tmpf+=tmp;
            tmpg+=g[v];
            if(mg.find(col[i])==mg.end())
                mg[col[i]]=g[v],mf[col[i]]=tmp;
            else mg[col[i]]+=g[v],mf[col[i]]+=tmp;
        }
    ans+=tmpf;
    tmp=0;
    for(int i=head[u],v;i>=0;i=next[i])
        if(!vis[v=ver[i]])
            tmp+=(tmpf-mf[col[i]])*g[v]+(tmpg-mg[col[i]])*f[v];
    ans+=tmp/2;
    vis[u]=0;
}
int main()
{
    int i,j,k,t,n;
    while(~scanf("%d",&n))
    {
        REP(i,n)scanf("%I64d",&a[i]);
        edge=0;
        REP(i,n)head[i]=-1,vis[i]=0;
        REP(t,n-1)
        {
            scanf("%d%d%d",&i,&j,&k);
            --i,--j;
            ver[edge]=j,col[edge]=k,next[edge]=head[i],head[i]=edge++;
            ver[edge]=i,col[edge]=k,next[edge]=head[j],head[j]=edge++;
        }
        ans=0;
        dp(0,0);
        printf("%I64d\n",ans);
    }
	return 0;
}