bzoj3127/3697 [Usaco2013 Open]Yin and Yang

# include <stdio.h>
# include <string.h>
# include <algorithm>
// # include <bits/stdc++.h>

using namespace std;

typedef long long ll;
typedef long double ld;
typedef unsigned long long ull;
const int M = 5e5 + 10;
const int mod = 1e9+7;

# define RG register
# define ST static

int n; ll ans = 0;
int head[M], nxt[M], to[M], w[M], tot=0;
bool vis[M]; 

inline void add(int u, int v, int _w) {
    ++tot; nxt[tot] = head[u]; head[u] = tot; to[tot] = v; w[tot] = _w; 
}
inline void adde(int u, int v, int _w) {
    add(u, v, _w); add(v, u, _w);
}

int sz[M], mx[M]; 
inline void dfsSize(int x, int fa = 0) {
    sz[x] = 1; mx[x] = 0;
    for (int i=head[x]; i; i=nxt[i]) {
        if(to[i] == fa || vis[to[i]]) continue;
        dfsSize(to[i], x); 
        sz[x] += sz[to[i]];
        if(sz[to[i]] > mx[x]) mx[x] = sz[to[i]]; 
    }
}

int centre = 0, mi;
inline void dfsCentre(int x, int tp, int fa = 0) {
    if(sz[tp] - sz[x] > mx[x]) mx[x] = sz[tp] - sz[x]; 
    if(mx[x] < mi) mi = mx[x], centre = x;
    for (int i=head[x]; i; i=nxt[i]) {
        if(to[i] == fa || vis[to[i]]) continue;
        dfsCentre(to[i], tp, x);
    }
}

int t[M]; 
ll f[M][2], g[M][2];
int mxd, vmin = 1e9, vmax = 0;
// g: cur tree, f: all
inline void dfsAns(int x, int v, int fa) { 
    if(t[v]) g[v][1] ++;
    else g[v][0] ++; 
    t[v] ++; 
    vmin = min(vmin, v), vmax = max(vmax, v); 
    for (int i=head[x]; i; i=nxt[i]) { 
        if(to[i] == fa || vis[to[i]]) continue;
        dfsAns(to[i], v+w[i], x);
    }
    t[v] --; 
}

inline void calcAns(int x) {
    int Vmin = 1e7, Vmax = 0; 
    f[n][0] = 1; 
    for (int i=head[x]; i; i=nxt[i]) {
        if(vis[to[i]]) continue; 
        vmin = 1e7, vmax = 0;
        dfsAns(to[i], n+w[i], x);
        Vmin = min(Vmin, vmin), Vmax = max(Vmax, vmax); 
        ans += (f[n][0]-1) * g[n][0]; //啥都不动，没有1，是没有意义的，为了计数方便，前面先赋值1. 
        for (int j=vmin-n; j<=vmax-n; ++j) 
            // 将g合并到f上，统计答案 
            ans += f[n-j][1] * g[n+j][1] + f[n-j][0] * g[n+j][1] + f[n-j][1] * g[n+j][0]; 
        for (int j=vmin; j<=vmax; ++j) {
            f[j][0] += g[j][0];
            f[j][1] += g[j][1];
            g[j][0] = g[j][1] = 0;
        }
    }
    for (int j=Vmin; j<=Vmax; ++j) f[j][0] = f[j][1] = 0; 
}
        

inline void dfs(int x) {
    dfsSize(x);
    mi = n; dfsCentre(x, x);
    // cross the centre
    calcAns(centre);
    // ================
    vis[centre] = 1; 
    for (int i=head[centre]; i; i=nxt[i]) {
        if(!vis[to[i]]) dfs(to[i]);
    }
}
 
int main() {
    scanf("%d", &n);
    for (int i=1; i<n; ++i) {
        int u, v, _w; scanf("%d%d%d", &u, &v, &_w); 
        adde(u, v, _w ? 1 : -1); 
    }
    dfs(1);
    printf("%lld\n", ans); 
    return 0;
}