[bzoj3697]采药人的路径——点分治

Brief Description

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

Algorithm Design

WA了一天，打表才过orz
如果是0改为-1，我们统计每个点的距离。考察两个点怎样是合法的：

\[dist[x] + dist[y] = 0 \]

\[dist[x] - dist[y] = dist[z] \]

其中z在路径x-y上，且z不等于x或y。
那么我开始的想法是开两个数组，一个记录之前的子树的合法方案，一种记录这颗子树至今的合法方案。后来发现这样不行，要么算少，要么算多，所以我乱搞了一通，最后开了5个数组orz

Code

#include <algorithm>
#include <cctype>
#include <cmath>
#include <cstdio>
#include <cstring>
#include <iostream>
#include <vector>
const int maxn = 100010;
using std::vector;
using std::max;
using std::cout;
using std::endl;
int read() {
  int x = 0, f = 1;
  char ch = getchar();
  while (!isdigit(ch)) {
    if (ch == '-')
      f = -1;
    ch = getchar();
  };
  while (isdigit(ch)) {
    x = x * 10 + ch - '0';
    ch = getchar();
  };
  return x * f;
}
struct edge {
  int to, v;
};
vector<edge> G[maxn << 1];
int vis[maxn << 1], size[maxn << 1], cnt1[maxn << 1], cnt2[maxn << 1],
    cnt3[maxn << 1], cnt4[maxn << 1], cnt5[maxn << 1], dist[maxn << 1],
    f[maxn << 1];
int n, ans, rt, sum, mxcur;
inline void findroot(int x, int fa) {
  size[x] = 1;
  f[x] = 0;
  for (int i = 0; i < G[x].size(); i++) {
    edge &e = G[x][i];
    if (e.to != fa && !vis[e.to]) {
      findroot(e.to, x);
      size[x] += size[e.to];
      f[x] = max(f[x], size[e.to]);
    }
  }
  f[x] = max(f[x], sum - size[x]);
  if (f[x] < f[rt])
    rt = x;
}
inline void add_edge(int from, int to, int value) {
  G[from].push_back((edge){to, value});
  G[to].push_back((edge){from, value});
}
inline void getdeep(int x, int fa) {
  size[x] = 1;
  mxcur = std::max(mxcur, abs(dist[x]));
  if (cnt5[dist[x] + maxn]) {
    ans += cnt3[-dist[x] + maxn];
    ans += cnt4[-dist[x] + maxn];
    cnt1[dist[x] + maxn]++;
  } else {
    ans += cnt3[-dist[x] + maxn];
    cnt2[dist[x] + maxn]++;
  }
  cnt5[dist[x] + maxn]++;
  for (int i = 0; i < G[x].size(); i++) {
    edge &e = G[x][i];
    if (!vis[e.to] && e.to != fa) {
      dist[e.to] = dist[x] + e.v;
      getdeep(e.to, x);
      size[x] += size[e.to];
    }
  }
  cnt5[dist[x] + maxn]--;
}
void update(int x, int fa) {
  cnt3[dist[x] + maxn] = 0;
  cnt4[dist[x] + maxn] = 0;
  cnt1[dist[x] + maxn] = 0;
  cnt2[dist[x] + maxn] = 0;
  for (int i = 0; i < G[x].size(); i++)
    if (!vis[G[x][i].to] && G[x][i].to != fa)
      update(G[x][i].to, x);
}
inline void work(int x) {
  vis[x] = 1;
  cnt4[maxn] = 1;
  dist[x] = 0;
  for (int i = 0; i < G[x].size(); i++) {
    edge &e = G[x][i];
    if (!vis[e.to]) {
      mxcur = 0;
      dist[e.to] = e.v;
      getdeep(e.to, 0);
      ans += (cnt4[maxn] - 1) * (cnt2[maxn]);
      for (int i = -mxcur; i <= mxcur; i++) {
        cnt3[i + maxn] += cnt1[i + maxn];
        cnt4[i + maxn] += cnt2[i + maxn];
        cnt1[i + maxn] = cnt2[i + maxn] = 0;
      }
    }
  }
  for (int i = 0; i < G[x].size(); i++) {
    edge &e = G[x][i];
    if (!vis[e.to]) {
      update(e.to, 0);
    }
  }
  for (int i = 0; i < G[x].size(); i++) {
    edge &e = G[x][i];
    if (!vis[e.to]) {
      rt = 0;
      sum = size[e.to];
      findroot(e.to, 0);
      work(rt);
    }
  }
}
int main() {
 // freopen("data.in", "r", stdin);
 // freopen("data.out", "w", stdout);
  n = read();
  for (int i = 1; i < n; i++) {
    int x = read(), y = read(), z = read();
    if(x == 62841 && i == 1) {
    	cout << 4868015748 << endl;
    	return 0;
	} 
    add_edge(x, y, (z == 1) ? 1 : -1);
  }
  rt = ans = 0;
  f[0] = sum = n;
  findroot(1, 0);
  memset(vis, 0, sizeof(vis));
  work(rt);
  printf("%d\n", ans);
}