题目地址：BZOJ：Problem 3697. — 采药人的路径

题目描述

采药人的药田是一个树状结构，每条路径上都种植着同种药材。
采药人以自己对药材独到的见解，对每种药材进行了分类。大致分为两类，一种是阴性的，一种是阳性的。
采药人每天都要进行采药活动。他选择的路径是很有讲究的，他认为阴阳平衡是很重要的，所以他走的一定是两种药材数目相等的路径。采药工作是很辛苦的，所以他希望他选出的路径中有一个可以作为休息站的节点（不包括起点和终点），满足起点到休息站和休息站到终点的路径也是阴阳平衡的。他想知道他一共可以选择多少种不同的路径。

题意简述

给你一棵树，树的边权为0或1，求满足下列条件的路径数：

1. 路径中的边一半边权为0，一半为1；
2. 路径中存在一个点把路径分为两段，这两段同样满足条件1。

输入输出格式

输入格式：
第1行包含一个整数N。
接下来N-1行，每行包含三个整数a_i、b_i和t_i，表示这条路上药材的类型。

输出格式：
输出符合采药人要求的路径数目。

输入输出样例

输入样例#1：

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

输出样例#1：

1

说明

对于100%的数据，N ≤ 100,000。

题解

参考资料：「BZOJ3697」「FJ2014集训」采药人的路径 – 点分治 – hzwer.com
看到这个题就能想到点分治，但是其实并不好办。
首先边权0和1考虑改成-1和1，这样做边权和比较方便。我们考虑一个重心，经过它的路径如何才算合法。只有两端到重心的路径上存在一个点到两端的路径上权值和为0。只要根到一个点的路径上存在一个点满足前面那个条件，它就可以和其他子树中的任意点构成一条路径。
我们考虑DFS处理子树的时候分开处理有无休息站的点，设f[d][0]为无休息站的，f[d][1]为有休息站的点的个数，另外开g[d][0/1]数组存前面子树的上述数组之和，发现每一个子树对答案的贡献是
\sum_i (f[i][1] \cdot g[i][1] + f[i][1] \cdot g[i][0] + f[i][0] \cdot g[i][1])
另外计算根为该休息站的情况的贡献是f[0][0] \cdot (g[0][0] - 1)。其中i的取值范围显然是[-最大深度, 最大深度]，这里可以加个n强行搞成非负数用来做下标。
但是需要注意，由于数组很大，memset会很慢，需要手动清零。

代码

// 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;
}