题目: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;
}