BZOJ3697 采药人的路径 【点分治】

题目

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

输入格式

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

输出格式

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

输入样例

7

1 2 0

3 1 1

2 4 0

5 2 0

6 3 1

5 7 1

输出样例

1

提示

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

题解

树上路径计数，联想到点分治

对于分治出来的，每一个子树的根，我们需要找出经过该根的满足条件的路径数
很容易想到，如果对两种权值赋值为-1和1的话，一条路径权值和为0一定是阴阳互补的
先不考虑休息站，如果求阴阳互补的路径数，只需要开一个数组记录各种权值出现的次数【当然由于有负数就加上个n】
对于根的每棵子树统计完时，先与之前统计的计算对答案的贡献，路经长为-x的个数乘以之前为x的个数即可

现在考虑休息站，意味着路径中出现中途已经出现0的情况，既然如此，那么休息站到根的权值和与该点到根的权值和一定相等
开一个数组记录dfs路径上各种权值出现的次数，看看该点到根是否存在这样可以作为休息站的点

这样所有的点就分为有休息站和没有休息站两种，分开组合统计，具体就自行思考吧【看代码也行】
复杂度\(O(nlogn)\)

#include<iostream>
#include<cmath>
#include<cstdio>
#include<cstring>
#include<algorithm>
#define LL long long int
#define REP(i,n) for (int i = 1; i <= (n); i++)
#define Redge(u) for (int k = h[u],to; k; k = ed[k].nxt)
#define BUG(s,n) for (int i = 1; i <= (n); i++) cout<<s[i]<<' '; puts("");
using namespace std;
const int maxn = 100005,maxm = 200005,INF = 1000000000;
inline int read(){
	int out = 0,flag = 1; char c = getchar();
	while (c < 48 || c > 57) {if (c == '-') flag = -1; c = getchar();}
	while (c >= 48 && c <= 57) {out = (out << 3) + (out << 1) + c - '0'; c = getchar();}
	return out * flag;
}
LL ans;
int n,h[maxn],ne = 2;
struct EDGE{int to,nxt,w;}ed[maxm];
inline void build(int u,int v,int w){
	ed[ne] = (EDGE){v,h[u],w}; h[u] = ne++;
	ed[ne] = (EDGE){u,h[v],w}; h[v] = ne++;
}
int sum,Siz[maxn],F[maxn],rt,vis[maxn];
void getRT(int u,int f){
	Siz[u] = 1; F[u] = 0;
	Redge(u) if (!vis[to = ed[k].to] && to != f){
		getRT(to,u);
		Siz[u] += Siz[to];
		F[u] = max(F[u],Siz[to]);
	}
	F[u] = max(F[u],sum - Siz[u]);
	if (F[u] < F[rt]) rt = u;
}
LL A[maxm],AA[maxm],B[maxm],BB[maxm];
int ex[maxm],d[maxn];
void dfs(int u,int f){
	int x = d[u] + n;
	if (ex[x]) BB[x]++;
	else B[x]++;
	ex[x]++;
	Redge(u) if (!vis[to = ed[k].to] && to != f){
		d[to] = d[u] + ed[k].w;
		dfs(to,u);
	}
	ex[x]--;
}
void DFS(int u,int f){
	Siz[u] = 1;
	Redge(u) if (!vis[to = ed[k].to] && to != f){
		DFS(to,u);
		Siz[u] += Siz[to];
	}
}
void solve(int u){
	vis[u] = true;
	DFS(u,0);
	Redge(u) if (!vis[to = ed[k].to]){
		for (int i = 0; i <= Siz[to]; i++)
			B[n + i] = B[n - i] = BB[n + i] = BB[n - i] = 0;
		d[to] = ed[k].w;
		dfs(to,u);
		ans += BB[n] + BB[n] * A[n] + B[n] * AA[n] + BB[n] * AA[n] + B[n] * A[n];
		for (int i = 1; i <= Siz[to]; i++){
			ans += BB[n - i] * A[n + i] + B[n - i] * AA[n + i] + BB[n - i] * AA[n + i];
			ans += BB[n + i] * A[n - i] + B[n + i] * AA[n - i] + BB[n + i] * AA[n - i];
		}
		A[n] += B[n]; AA[n] += BB[n];
		for (int i = 1; i <= Siz[to]; i++){
			A[n - i] += B[n - i];
			AA[n - i] += BB[n - i];
			A[n + i] += B[n + i];
			AA[n + i] += BB[n + i];
		}
	}
	for (int i = 0; i <= Siz[u]; i++)
		A[n - i] = A[n + i] = AA[n - i] = AA[n + i] = 0;
	Redge(u) if (!vis[to = ed[k].to]){
		sum = Siz[to]; F[rt = 0] = INF;
		getRT(to,0);
		solve(rt);
	}
}
int main(){
	n = read(); int a,b,w;
	for (int i = 1; i < n; i++){
		a = read(); b = read(); w = read();
		build(a,b,w ? 1 : -1);
	}
	sum = n; F[rt = 0] = INF;
	getRT(1,0);
	solve(rt);
	cout << ans << endl;
	return 0;
}