【HDU4303】Hourai Jeweled
题意
有一棵n个结点的树，每个结点都有一个值，没一条边都有一个颜色。如果某条路径上，相邻的边颜色不同，那么把这路径上所有的点的值加起来。 输出所有符合条件的路径上值的和。 n<=300000。
分析
场上读错题了。当时以为是路径上所有的边都不能有相同的颜色，雾。。。
树形DP。
1 #include <cstdio> 2 #include <cstring> 3 #include <iostream> 4 #include <algorithm> 5 #include <map> 6 using namespace std; 7 const int maxn=300000+10; 8 int head[maxn],Next[2*maxn],to[2*maxn],col[2*maxn],val[maxn]; 9 int n,sz; 10 long long ans; 11 void add_edge(int a,int b,int c){ 12 sz++; 13 to[sz]=b; 14 col[sz]=c; 15 Next[sz]=head[a]; 16 head[a]=sz; 17 } 18 long long f[maxn],num[maxn]; 19 void dp(int u,int fa,int co){ 20 long long alltmp=0,allnum=0; 21 num[u]=1;f[u]=val[u]; 22 map<int,long long>Mdp; 23 map<int,long long>Mnum; 24 25 for(int i=head[u];i;i=Next[i]){ 26 int v=to[i]; 27 if(v!=fa){ 28 dp(v,u,col[i]); 29 if(co!=col[i]){ 30 f[u]+=f[v]+val[u]*num[v]; 31 num[u]+=num[v]; 32 } 33 ans+=f[v]+val[u]*num[v]; 34 alltmp+=f[v]+val[u]*num[v]; 35 allnum+=num[v]; 36 Mnum[col[i]]+=num[v]; 37 Mdp[col[i]]+=f[v]+val[u]*num[v]; 38 } 39 } 40 long long temp=0; 41 for(int i=head[u];i;i=Next[i]){ 42 if(to[i]!=fa){ 43 temp+=f[to[i]]*(allnum-Mnum[col[i]])+num[to[i]]*(alltmp-Mdp[col[i]]); 44 } 45 } 46 ans+=temp/2; 47 } 48 int main(){ 49 while(scanf("%d",&n)!=EOF){ 50 ans=0; 51 memset(head,0,sizeof(head)); 52 sz=0; 53 for(int i=1;i<=n;i++) 54 scanf("%d",&val[i]); 55 int a,b,c; 56 for(int i=1;i<n;i++){ 57 scanf("%d%d%d",&a,&b,&c); 58 add_edge(a,b,c); 59 add_edge(b,a,c); 60 } 61 dp(1,0,0); 62 printf("%lld\n",ans); 63 } 64 return 0; 65 }
