Hdu 4303 Hourai Jeweled

Hourai Jeweled

官方題解：
從任意一點開始深搜，每顆子樹搜索完畢之後向上返回pair<可以延伸到該點且最後一條邊與由父節點到該點的邊顏色不同的gorgeous邊的條數 , 所有這種邊分數的總和>
每次深搜完一個子節點之後，增加的過這一點的gorgeous邊的總分數爲：
    之前深搜的所有子節點向上返回的邊數之和 * 當前子節點返回的分數 +
    之前深搜的所有子節點向上返回的分數之和 * 當前子節點返回的邊數 +
    之前深搜的所有子節點向上返回的邊數之和 * 當前子節點返回的邊數 * 當前點的權

用O(n)的時間把所有的累計起來就好了什麼的

纔沒有呢 = =
如果有當前節點分別到兩個子節點的邊的顏色相同的話也是不行的,
由於數據中添加了度數很大的點，理論上O(兒子數^2)的兩兩統計也是要被卡掉的 （希望確實的做到了
正確的做法是對所有的兒子按邊的顏色排個序，然後在按這個序深搜和統計

其實就是統計出每個點被使用了多少次。我們可以邊向上推，邊算出答案。對於節點j，它有如下兩種情況
1.做爲終點，包括從子孫節點到該點的路徑以及從該點延伸到祖先節點,
2.做爲中間節點，包括從子孫節點x到子孫節點y和從子孫節點x到祖先節點p.
對於做爲終點的情況從子孫節點延伸到該點的路徑的，我們可以直接dfs獲得
對於做爲中間節點：從孩子節點x走到孩子節點y的情況，我們可以在dfsj的兒子節點的時候獲得

對於從該點向上延伸的情況，我們直接返回給父親節點。具體的可以參看程序理解。

/*
    author    : csuchenan
    prog      : hdu4303
    algorithm : Tree DP
    csu_chenan	562MS	13232K	2874B	C++	2012-10-25 14:24:32
*/
#include <cstdio>
#include <cstring>
#include <vector>
#include <utility>
#include <algorithm>
using std::sort;
using std::pair;
using std::vector;
using std::make_pair;
typedef pair<int, int> Edge;
typedef long long LL;
const int maxn = 300000;
vector<Edge> G[maxn+5];
int val[maxn+5];
int n;
LL ans;

void init(){
    for(int i = 0; i <= n + 2 ; i ++){
        G[i].clear();
    }
}

bool read(){
    init();
    if(scanf("%d", &n)==EOF)
        return false;
    for(int i = 1; i <= n; i ++){
        scanf("%d", &val[i]);
    }

    int tv, tu, tc;
    for(int i = 1; i < n; i ++){
        scanf("%d%d%d", &tv, &tu, &tc);
        G[tv].push_back(make_pair(tu, tc));
        G[tu].push_back(make_pair(tv, tc));
    }
    return true;
}

//將每個點相鄰的邊按照顏色排序
bool cmp(Edge a, Edge b){
    return a.second < b.second;
}

//v當前節點，f爲父親節點，score爲向根傳遞的以v爲終點的總分
//num爲返回的以v終點的路徑的數目，cc爲(f,v)的顏色
void dfs(int v, int f, int cc, LL &score, LL &num){
    score = 0;
    num = 0;

    int lastc = -1;//當前兒子節點的前一個兄弟節點的顏色
    LL lscore = 0;
    LL tscore = 0; //當前兒子節點之前的分數總和
    LL tnum = 0;   //當前兒子節點之前的邊數總和
    LL lnum = 0;

    sort(G[v].begin(), G[v].end(), cmp);
    for(int i = 0; i < G[v].size(); i ++){
        int u = G[v][i].first;
        int c = G[v][i].second;
        //如果u爲父親節點
        if(u==f){
            continue;
        }
        LL ss, ne; //ss是當前兒子節點訪問得到的分數總和，ne是節點個數
        //如果是兒子節點
        dfs(u, v, c, ss, ne);
        if(cc != c){
            score += ss + val[v] * ne;
            num += ne;
        }
        //將以v爲終點的節點加到答案中
        ans += ne * val[v] + ss;
        //當前兒子節點的的邊顏色與兄弟節點的顏色相同
        if(c==lastc){
            ans += lnum * ss;
            ans += lscore * ne;
            ans += val[v] * lnum * ne;
        }
        else{
            ans += tnum * ss;
            ans += tscore * ne;
            ans += val[v] * tnum * ne;
        }
        //這裏更新上一次訪問的邊數的和時需要注意
        if(lastc != -1 && c != lastc){
            lnum = tnum;
            lscore = tscore;
        }
        tnum += ne;
        tscore += ss;
        lastc = c;
    }
    //對於v點也要包含在內
    if(f != 0){
        score += val[v];
        num ++;
    }
}

void solve(){
    LL score = 0 , num  = 0;
    ans = 0;
    dfs(1, 0, -1, score, num);
    printf("%I64d\n", ans);
}

int main(){
    while(read()){
        solve();
    }
    return 0;
}