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
, w[M], cnt = 1, node, rt;
long long ans = 0;
int d
, sz
;
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。