bzoj3697 采药人的路径

Description

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

Input

第 1 行包含一个整数 N 。
接下来 N−1 行，每行包含三个整数 ai,bi,ti，表示这条路上药材的类型。

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≤100000 。

Solution

本周点分治最后一刷！
如何判断一条从根出发的路径是否包含休息站？只要在 dfs 中记录下这条路径前缀和 x ，同时用一个桶判断这条路径是否存在前缀和为 x 的节点。
用 f[i][0...1] 表示当前子树不存在/存在休息站的长度为 i 路径数目。 g[i][0...1] 表示之前访问过的子树不存在/存在休息站的长度为 i 路径数目。
那么当前子树的贡献就是 f[0][0]∗g[0][0]+Σf[i][0]∗g[−i][1]+f[i][1]∗g[−i][0]+f[i][1]∗g[−i][1]。
其中 i∈[−dep,dep] ， dep 为当前子树的最大深度。
另外我写的时候在 getRoot 的时候写错了一些东西，于是

#include<bits/stdc++.h>using namespace std;#define N 200005#define rep(i, a, b) for (int i = a; i <= b; i++)#define drp(i, a, b) for (int i = a; i >= b; i--)#define fech(i, x) for (int i = 0; i < x.size(); i++)#define ll long longinline int read() {    int x = 0, flag = 1; char ch = getchar(); while (!isdigit(ch)) { if (!(ch ^ '-')) flag = -1; ch = getchar(); }    while (isdigit(ch)) x = (x << 1) + (x << 3) + ch - '0', ch = getchar(); return x * flag;}inline void write(int x) {    if (!x) { putchar('0'); return; } if (x < 0) putchar('-'), x = -x;    char buf[20] = ""; int top = 0; while (x) buf[++top] = x % 10 + '0', x /= 10; while (top) putchar(buf[top--]);}int n;struct edgeType { int u, v, w; }eg[N]; int tot;#define getEg edgeType e = eg[g[u][i]]vector<int> g[N];bool vis[N];int Size[N], mx[N], sum, root, mxDep;int dis[N], dep[N], t[N];ll ans, F[N][2], G[N][2];void getRoot(int u, int fa) {    Size[u] = 1; mx[u] = 0;    fech(i, g[u]) {        getEg; if (!(e.v ^ fa) || vis[e.v]) continue;        getRoot(e.v, u); Size[u] += Size[e.v], mx[u] = max(mx[u], Size[e.v]);    }    mx[u] = max(mx[u], sum - Size[u]);    if (mx[root] > mx[u]) root = u;}void calcDis(int u, int fa) {    mxDep = max(mxDep, dep[u]);    F[dis[u]][(bool)t[dis[u]]]++; t[dis[u]]++;    fech(i, g[u]) {        getEg; if (!(e.v ^ fa) || vis[e.v]) continue;        dis[e.v] = dis[u] + e.w, dep[e.v] = dep[u] + 1; calcDis(e.v, u);    }    t[dis[u]]--;}void solve(int u) {    int mx = 0; vis[u] = 1; G[n][0] = 1;    fech(i, g[u]) {        getEg; if (vis[e.v]) continue;        dep[e.v] = 1; dis[e.v] = e.w + n; mxDep = 1; calcDis(e.v, 0); mx = max(mx, mxDep);        ans += (G[n][0] - 1) * F[n][0];        rep(j, -mxDep, mxDep)            ans += G[n - j][1] * F[n + j][1] + G[n - j][0] * F[n + j][1] + G[n - j][1] * F[n + j][0];        rep(j, n - mxDep, n + mxDep) rep(k, 0, 1) G[j][k] += F[j][k], F[j][k] = 0;    }    rep(i, n - mx, n + mx) rep(j, 0, 1) G[i][j] = 0;    fech(i, g[u]) {        getEg; if (vis[e.v]) continue;        sum = Size[e.v]; root = 0; getRoot(e.v, 0); solve(root);    }}int main() {    n = read(); rep(i, 2, n) {        int u = read(), v = read(), w = read(); if (!w) w = -1;        eg[++tot] = edgeType{ u, v, w }; g[u].push_back(tot);        eg[++tot] = edgeType{ v, u, w }; g[v].push_back(tot);    }    sum = mx[0] = n; getRoot(1, 0); solve(root); printf("%lld", ans);    return 0;}