Hourai Jeweled

HDU - 4303
Hourai Jeweled

Time Limit: 2000MS Memory Limit: 163840KB 64bit IO Format: %I64d & %I64u

Description

Kaguya Houraisan was once a princess of the Lunarians, a race of people living on the Moon. She was exiled to Earth over a thousand years ago for the crime of using the forbidden Hourai Elixir to make herself immortal. Tales of her unearthly beauty led men from all across the land to seek her hand in marriage, but none could successfully complete her trial of the Five Impossible Requests.

One of these requests is to reckon the value of “Hourai Jeweled (蓬莱の玉の枝)”. The only one real treasure Kaguya has, in her possession. As showed in the picture, Hourai Jeweled is a tree-shaped twig. In which, each node is ornamented with a valuable diamond and each edge is painted with a briliant color (only bright man can distinguish the difference). Due to lunarians’ eccentric taste, the value of this treasure is calculated as all the gorgeous roads’ value it has. The road between two different nodes is said to be gorgeous, if and only if all the adjacent edges in this road has diffenrent color. And the value of this road is the sum of all the nodes’ through the road.
Given the value of each node and the color of each edge. Could you tell Kaguya the value of her Hourai Jeweled?

Input

The input consists of several test cases.
The first line of each case contains one integer $N (1 <= N <= 300000)$, which is the number of nodes in Hourai Jeweled.
The second line contains $N$ integers, the $i$-th of which is $Vi (1 <= Vi <= 100000)$, the value of node $i$.
Each of the next $N-1$ lines contains three space-separated integer $X$, $Y$and $Z (1<=X,Y<=N, 1 <= Z <= 100000)$, which represent that there is an edge between $X$ and $Y$ painted with colour $Z$.

Output

For each test case, output a line containing the value of Hourai Jeweled.

Sample Input

6
6 2 3 7 1 4
1 2 1
1 3 2
1 4 3
2 5 1
2 6 2

Sample Output

134

Solution & Code

考虑用树形dp解决。
讨论到每个点$u$，其儿子为$v$，父亲为$fa$共有三种路径（如果合法）可以贡献答案：
1、$...-v-u$
2、$...-v-u-fa-...$
3、$...-v_1-u-v_2-...$
对每个点记录两个状态：
1、$g[u]$表示情况2中$u$及其子树贡献的权值和（每个点的权值可能被加多次）
2、$h[u]$表示情况2中通过$u$的路径数
除此之外还可以把边的颜色转换成点的颜色，用$clr[u]$表示$u-fa$边的颜色
则$g$与$h$的状态转移如下：
1、$h[u]=\Sigma h[v](clr[u]==clr[v])+1$
2、$g[u]=\Sigma g[v](clr[u]==clr[v])+val[u]*h[u]$
按之前的三种情况分别统计答案：
1、$ans+=\Sigma g[v]+val[u]*h[u]$
2、因为路径只确定了一个端点，所以对答案没有贡献
3、$ans+=\Sigma (g[v_1]*h[v_2]+g[v_2]*h[v_1]+val[u]*h[v_1]*h[v_2])/2(clr[v_1]==clr[v_2])$
暴力统计情况3是会超时的，所以可以记录每种颜色的儿子的$g[v]$和，和所有颜色的儿子的$g[v]$和，对数组也是同样处理，详见代码

#include <cstdio>
#include <cstring>
#include <algorithm>
using namespace std;

typedef long long ll;

const int maxn = 300005;
const int maxm = 600005;

ll ans, g[maxn], h[maxn], val[maxn];
int n, clr[maxn];

int tot;
int lst[maxn], nxt[maxm], nd[maxm], c[maxm];
void add_edge(int u, int v, int t){

    ++tot; nxt[tot] = lst[u]; nd[tot] = v; c[tot] = t; lst[u] = tot;
    ++tot; nxt[tot] = lst[v]; nd[tot] = u; c[tot] = t; lst[v] = tot;
}

ll g_clr[maxn], h_clr[maxn]; // 每种颜色的g[v]和与h[v]和
void dfs(int u, int fa){

    ll g_sum = 0, h_sum = 0; //所有颜色的g[v]和与h[v]和
    for(int l = lst[u]; l; l = nxt[l]){
        int v = nd[l];
        if(v == fa) continue;
        clr[v] = c[l]; // 将边的颜色转化为点的颜色
        dfs(v, u);
    }
    for(int l = lst[u]; l; l = nxt[l]) g_clr[clr[nd[l]]] = h_clr[clr[nd[l]]] = 0; // 讨论完儿子后将有用的部分清零
    for(int l = lst[u]; l; l = nxt[l]){
        int v = nd[l];
        if(v == fa) continue;
        g_clr[clr[v]] += g[v];
        g_sum += g[v];
        h_clr[clr[v]] += h[v];
        h_sum += h[v];
    }
    h[u] = 1;
    g[u] = val[u];
    ll tmp = 0;
    for(int l = lst[u]; l; l = nxt[l]){
        int v = nd[l];
        if(v == fa) continue;
        ans += g[v] + val[u] * h[v]; // 情况1
        tmp += g[v] * (h_sum - h_clr[clr[v]]) + // 情况3 
               h[v] * (g_sum - g_clr[clr[v]]) + 
               val[u] * h[v] * (h_sum - h_clr[clr[v]]);
        if(clr[v] != clr[u]){ // 转移
            h[u] += h[v];
            g[u] += g[v] + val[u] * h[v];
        }
    }
    tmp /= 2; // 情况三种每条边都被算了2次
    ans += tmp;
}

void work(){

    ans = tot = 0;
    memset(lst, 0, sizeof lst);

    for(int i = 1; i <= n; ++i) scanf("%I64d", &val[i]);
    for(int i = 1; i < n; ++i){
        int u, v, t;
        scanf("%d%d%d", &u, &v, &t);
        add_edge(u, v, t);
    }
    dfs(1, 0);
    printf("%I64d\n", ans);

}