BZOJ 3697 采药人的路径 点分治

由于不知道两种颜色的和，但是如果将两种颜色分别设为1和-1，那么符合条件的和必为0，以这个树分治即可。
数组必须可以索引负值为了不调试使代码长了好多。。

#include <queue>
#include <cstdio>
#include <algorithm>
const int N = 100005, M = N * 2;
using namespace std;
struct SimpleArray {
    int c[2];
    int &operator [](int x) { return c[x]; }
};
template<class T>
struct Array {
    T c[N * 2];
    T &operator [](int x) { return c[N + x]; }
};
int vis[M], p[M], v[M], h[N], w[M], cnt = 1, node, rt;
long long ans = 0;
int d[N], sz[N];
Array<SimpleArray> f, g;
Array<int> s;
void add(int a, int b, int c) {
    p[++cnt] = h[a]; v[cnt] = b; w[cnt] = c; h[a] = cnt;
}
void root(int x, int fa) {
    sz[x] = 1; d[x] = 0;
    for (int i = h[x]; i; i = p[i])
        if (v[i] != fa && !vis[i]) {
            root(v[i], x);
            sz[x] += sz[v[i]];
            d[x] = max(d[x], sz[v[i]]);
        }
    d[x] = max(d[x], node - sz[x]);
    if (d[x] < d[rt]) rt = x;
}
int get_root(int x, int fa, int sz) {
    rt = 0; node = sz; d[0] = 2147483647;
    root(x, fa); return rt;
}
void dfs_seq(int x, int fa, int dep, int dis, int &mxdep) {
    mxdep = max(mxdep, dep);
    ++g[dis][s[dis] ? 1 : 0];
    ++s[dis];
    for (int i = h[x]; i; i = p[i])
        if (!vis[i] && v[i] != fa)
            dfs_seq(v[i], x, dep + 1, dis + w[i], mxdep);
    --s[dis];
}
void work(int x) {
    int i, j, mxdep, mx = 0;
    f[0][0] = 1;
    for (i = h[x]; i; i = p[i])
        if (!vis[i]) {
            mxdep = 1; dfs_seq(v[i], x, 1, w[i], mxdep);
            mx = max(mx, mxdep);
            ans += g[0][0] * (f[0][0] - 1);
            for (j = -mxdep; j <= mxdep; ++j)
                ans += g[j][0] * f[-j][1] + g[j][1] * f[-j][0] + g[j][1] * f[-j][1];
            for (j = -mxdep; j <= mxdep; ++j) {
                f[j][0] += g[j][0];
                f[j][1] += g[j][1];
                g[j][0] = g[j][1] = 0;
            }
        }
    for (j = -mx; j <= mx; ++j)
        f[j][0] = f[j][1] = 0;
    for (i = h[x]; i; i = p[i])
        if (!vis[i]) {
            vis[i] = vis[i ^ 1] = true;
            work(get_root(v[i], x, sz[v[i]]));
        }
}
int main() {
    int a, b, c, i, n;
    scanf("%d", &n);
    for (i = 1; i < n; ++i) {
        scanf("%d%d%d", &a, &b, &c);
        c = c ? 1 : -1;
        add(a, b, c); add(b, a, c);
    }
    work(get_root(1, 0, n));
    printf("%lld", ans);
    return 0;
}

3697: 采药人的路径

Description

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

Input

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

Output

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

Sample Input

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

Sample Output

1

HINT

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