【树形DP】HDU 4303 Hourai Jeweled

#include <cstdio>
#include <cstring>
#include <algorithm>

using namespace std;

const int MAX_N = 300007;

typedef long long ll;

struct Edge {
    int u, v, w, nxt;
    Edge () {
        
    }
    Edge (int _u, int _v ,int _w, int _n) {
        u = _u; v = _v; w = _w; nxt = _n;
    }
    bool operator < (const Edge &rhs) const {
        return w < rhs.w;
    }
};

int n;
int head[MAX_N], ecnt;
ll ans, val[MAX_N], num[MAX_N], sum[MAX_N];
Edge T[MAX_N], G[MAX_N << 1];

void Clear() {
    ecnt = ans = 0; memset(head, -1, sizeof head);
    memset(num, 0, sizeof num); memset(sum, 0, sizeof sum);
}

void add(int u, int v, int w) {
    G[ecnt] = Edge(u, v, w, head[u]);
    head[u] = ecnt++;
}

void dfs(int u, int fa, int col) {
    ll tot = 0, tot1 = 0, tot2 = 0, k = 0, k1 = 0, k2 = 0;
    int nowCol = 0;
    for (int i = head[u]; ~i; i = G[i].nxt) {
        int v = G[i].v, w = G[i].w;
        if (v == fa) continue;
        dfs(v, u, w);
        if (w != col) tot += sum[v], k += num[v];
        if (w != nowCol) {
            tot1 += tot2; k1 += k2;
            tot2 = sum[v], k2 = num[v];
            nowCol = w;
        } else tot2 += sum[v], k2 += num[v];
        ans += sum[v] * (k1 + 1) + tot1 * num[v] + val[u] * (k1 + 1) * num[v];
    }
    num[u] = ++k, sum[u] = val[u] * k + tot;
}

template <class T>
inline bool rd(T &ret) {
    char c; int sgn;
    if(c = getchar() , c == EOF) return false;
    while(c != ‘-‘ && (c < ‘0‘ || c > ‘9‘)) c = getchar();
    sgn = (c == ‘-‘) ? -1 : 1;
    ret = (c == ‘-‘) ? 0 : (c - ‘0‘);
    while(c = getchar(), c >= ‘0‘ && c <= ‘9‘) ret = ret * 10 + (c - ‘0‘);
    ret *= sgn;
    return true;
}

int main() {
    while (1 == scanf("%d", &n)) {
        for (int i = 1; i <= n; ++i) rd(val[i]);
        for (int i = 0; i + 1 < n; ++i) {
            rd(T[i].u), rd(T[i].v), rd(T[i].w);
        }
        sort(T, T + n - 1);
        Clear();
        for (int i = 0; i + 1 < n; ++i) {
            add(T[i].u, T[i].v, T[i].w);
            add(T[i].v, T[i].u, T[i].w);
        }
        dfs(1, -1, -1);
        printf("%I64d\n", ans);
    }
    return 0;
}