题目地址：洛谷：【P3085】[USACO13OPEN]阴和阳Yin and Yang – 洛谷、BZOJ：Problem 3127. — [Usaco2013 Open]Yin and Yang

题目描述

Farmer John is planning his morning walk on the farm. The farm is structured like a tree: it has N barns (1 <= N <= 100,000) which are connected by N-1 edges such that he can reach any barn from any other. Farmer John wants to choose a path which starts and ends at two different barns, such that he does not traverse any edge twice. He worries that his path might be a little long, so he also wants to choose another “rest stop” barn located on this path (which is distinct from the start or the end).
Along each edge is a herd of cows, either of the Charcolais (white hair) or the Angus (black hair) variety. Being the wise man that he is, Farmer John wants to balance the forces of yin and yang that weigh upon his walk. To do so, he wishes to choose a path such that he will pass by an equal number of Charcolais herds and Angus herds– both on the way from the start to his rest stop, and on the way from the rest stop to the end.
Farmer John is curious how many different paths he can choose that are “balanced” as described above. Two paths are different only if they consist of different sets of edges; a path should be counted only once even if there are multiple valid “rest stop” locations along the path that make it balanced.
Please help determine the number of paths Farmer John can choose!
给出一棵n个点的树，每条边为黑色或白色。问满足以下条件的路径条数：路径上存在一个不是端点的点，使得两端点到该点的路径上两种颜色的边数都相等。

输入输出格式

输入格式：
* Line 1: The integer N.
* Lines 2..N: Three integers a_i, b_i and t_i, representing the two barns that edge i connects. t_i is 0 if the herd along that edge is Charcolais, and 1 if the herd is Angus.

输出格式：
* Line 1: One integer, representing the number of possible paths Farmer John can choose from.

输入输出样例

输入样例#1：

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

输出样例#1：

1 

说明

There are 7 barns and 6 edges. The edges from 1 to 2, 2 to 4 and 2 to 5 have Charcolais herds along them.
No path of length 2 can have a suitable rest stop on it, so we can only consider paths of length 4. The only path that has a suitable rest stop is 3-1-2-5-7, with a rest stop at 2.

题解

同[BZOJ3697]采药人的路径一题。

代码

// Code by KSkun, 2018/5
#include <cstdio>
#include <cctype>
#include <cstring>

#include <algorithm>

typedef long long LL;

inline char fgc() {
    static char buf[100000], *p1 = buf, *p2 = buf;
    return p1 == p2 && (p2 = (p1 = buf) + fread(buf, 1, 100000, stdin), p1 == p2) 
        ? EOF : *p1++;
}

inline LL readint() {
    register LL res = 0, neg = 1;
    register char c = fgc();
    while(!isdigit(c)) {
        if(c == '-') neg = -1;
        c = fgc();
    }
    while(isdigit(c)) {
        res = (res << 1) + (res << 3) + c - '0';
        c = fgc();
    }
    return res * neg;
}

const int MAXN = 200005, INF = 1e9;

struct Edge {
    int to, type, nxt;
} gra[MAXN << 1];
int head[MAXN], tot;

inline void addedge(int u, int v, int type) {
    gra[tot] = Edge {v, type, head[u]}; head[u] = tot++;
    gra[tot] = Edge {u, type, head[v]}; head[v] = tot++;
}

int n;
LL ans;

int siz[MAXN], rt, rtsiz;
bool vis[MAXN];

inline void findrt(int u, int f, int tot) {
    siz[u] = 1; int mxsiz = 0;
    for(int i = head[u]; ~i; i = gra[i].nxt) {
        int v = gra[i].to, t = gra[i].type;
        if(v == f || vis[v]) continue;
        findrt(v, u, tot);
        siz[u] += siz[v];
        mxsiz = std::max(mxsiz, siz[v]);
    }
    mxsiz = std::max(mxsiz, tot - siz[u]);
    if(mxsiz < rtsiz) {
        rt = u; rtsiz = mxsiz;
    }
}

int f[MAXN][2], g[MAXN][2], mxdep, cnt[MAXN];

inline void calsubt(int u, int fa, int d, int dep) {
    mxdep = std::max(mxdep, dep);
    if(cnt[d]) f[d][1]++; else f[d][0]++;
    cnt[d]++;
    for(int i = head[u]; ~i; i = gra[i].nxt) {
        int v = gra[i].to, t = gra[i].type;
        if(v == fa || vis[v]) continue;
        calsubt(v, u, d + gra[i].type, dep + 1);
    }
    cnt[d]--;
}

inline void work(int u) {
    int mx = 0;
    g[n][0] = 1;
    for(int i = head[u]; ~i; i = gra[i].nxt) {
        int v = gra[i].to, t = gra[i].type;
        if(vis[v]) continue;
        mxdep = 1; calsubt(v, u, n + t, 1);
        mx = std::max(mx, mxdep);
        ans += 1ll * (g[n][0] - 1) * f[n][0];
        for(int j = -mxdep; j <= mxdep; j++) {
            ans += 1ll * g[n - j][1] * f[n + j][1] + 1ll * g[n - j][1] * f[n + j][0] 
                + 1ll * g[n - j][0] * f[n + j][1];
        }
        for(int j = -mxdep; j <= mxdep; j++) {
            g[n + j][0] += f[n + j][0];
            g[n + j][1] += f[n + j][1];
            f[n + j][0] = f[n + j][1] = 0;
        }
    }
    for(int i = -mx; i <= mx; i++) {
        g[n + i][0] = g[n + i][1] = 0;
    }
}

inline void divide(int u) {
    vis[u] = true;
    work(u);
    for(int i = head[u]; ~i; i = gra[i].nxt) {
        int v = gra[i].to, t = gra[i].type;
        if(vis[v]) continue;
        rt = 0; rtsiz = INF; findrt(v, u, siz[v]);
        divide(rt);
    }
}

int u, v, t;

int main() {
    memset(head, -1, sizeof(head));
    n = readint();
    for(int i = 1; i < n; i++) {
        u = readint(); v = readint(); t = readint();
        addedge(u, v, t ? t : -1);
    }
    rt = -1; rtsiz = INF; findrt(1, 0, n);
    divide(rt);
    printf("%lld", ans);
    return 0;
}