hdu4303 Hourai Jeweled 树形dp，统计

hdu4303 树形dp，统计
给一个树，边有颜色，点有权值。
满足某种条件的路径是好的路径，求(好的路径的(路径上的点的权值和)的总和)。
这种条件就是这个路径上没有两个相邻边同色。

关键就是算一个点的权值要被计算几次。
其实情况就两种，不要重复计算就可以。
情况1.包含(u,fa)这条边的路径。
情况2.不经过fa，而包含(u,v)这条边的路径。
这里要注意，如果有多个儿子节点，(v2,u,v1),(v1,u,v2)不要被重复计算。

每个点的儿子按颜色排序不会超过O(nlogn)复杂度。

我的代码居然也有看起来比较短的时候。。

#include<iostream>
#include<cstdio>
#include<cstring>
#include<set>
#include<vector>
#include<algorithm>
using namespace std;
#define pb push_back
#define NN 301000

typedef pair<int,int> pii;
vector<pii> ed[NN];

long long ans;
int son[NN],fc[NN];
int val[NN];

void dfs2(int u,int fa,int ecol,int fcnt){
    ans+=((long long)son[u]-fc[u])*fcnt*val[u];//情况1
    int i,j,sz=ed[u].size(),v,tcol=-1,vv,tfcnt;
    int totc=0,tott=0;
    for(i=0,j=0;i<sz;++i){
        v=ed[u][i].second;
        if (v==fa) continue;
        if (tcol!=ed[u][i].first){
            tcol=ed[u][i].first;
            totc=0;
            if (j<sz&&ed[u][j].second==fa) ++j;
            for(;j<sz&&ed[u][j].first==tcol;++j){
                vv=ed[u][j].second;
                if (vv==fa) continue;
                totc=totc+son[vv]-fc[vv];
            }
            tott+=totc;
        }
        ans+=((long long)son[u]-tott)*(son[v]-fc[v])*val[u];//情况2
        if (ecol==ed[u][i].first) tfcnt=son[u]-totc;
        else tfcnt=son[u]-totc+fcnt;
        dfs2(v,u,ed[u][i].first,tfcnt);
    }
}

int dfs1(int u,int fa,int ecol){
    son[u]=1;
    fc[u]=0;
    int i,sz=ed[u].size(),v,tmp;
    sort(ed[u].begin(),ed[u].end());
    for(i=0;i<sz;++i){
        if (ed[u][i].second==fa) continue;
        v=ed[u][i].second;
        tmp=dfs1(v,u,ed[u][i].first);
        son[u]+=tmp;
        if (ecol==ed[u][i].first) fc[u]+=tmp;
    }
    return son[u]-fc[u];
}

int main(){
    //freopen("4303in.txt","r",stdin);
    int n,i,a,b,c;
    while(scanf("%d",&n)!=EOF){
        for(i=1;i<=n;++i){
            scanf("%d",&val[i]);
            ed[i].clear();
        }
        for(i=1;i<n;++i){
            scanf("%d%d%d",&a,&b,&c);
            ed[a].pb(make_pair(c,b));
            ed[b].pb(make_pair(c,a));
        }
        dfs1(1,-1,-1);
        ans=0;
        dfs2(1,-1,-1,0);
        printf("%I64d\n",ans);
    }
    return 0;
}