[USACO13OPEN]Yin and Yang

我们考虑使用点分治，先找到一条路径，然后判断它是否可行？

经过研究，我们发现一条路径是可行解，当且仅当它是多条平衡路径合并而来的，于是我们把以根节点出发的路径分为两种，一种是包含有平衡路径的路径，它的贡献就是在其他子树中能够与它组成平衡路径的路径数目。

另一种是不包含有平衡路径的路径，则它的贡献就是其他子树的路径中能与它组成平衡路径且包含有平衡路径的路径数。

最后特判一下以根节点为端点的可行解。

以上这些路径都可以用map维护快速求解。

细节见代码：

#include <map>
#include <cstdio>
#include <algorithm>
using namespace std;
typedef long long LL;
void read(int &x){ 
	int f=1;x=0;char s=getchar();
	while(s<'0'||s>'9'){if(s=='-')f=-1;s=getchar();}
	while(s>='0'&&s<='9'){x=x*10+s-'0';s=getchar();}
	x*=f;
}
const int MAXN = 1e5 + 5;
int head[MAXN] , to[MAXN << 1] , nxt[MAXN << 1] , cnt , edge[MAXN << 1];
void add(int u , int v , int w) {nxt[++cnt] = head[u];head[u] = cnt;to[cnt] = v;edge[cnt] = w;}
int n , siz[MAXN] , rt , _max , vis[MAXN];
LL ans;
void get_siz(int x , int fa) {
	siz[x] = 1;
	for (int i = head[x]; i; i = nxt[i]) {
		if(to[i] == fa || vis[to[i]]) continue;
		get_siz(to[i] , x);
		siz[x] += siz[to[i]];
	}
}
void get_root(int x , int fa , int num) {
	int res = 0;
	for (int i = head[x]; i; i = nxt[i]) {
		if(to[i] == fa || vis[to[i]]) continue;
		get_root(to[i] , x , num);
		if(siz[to[i]] > res) res = siz[to[i]];
	}
	if(num - siz[x] > res) res = num - siz[x];
	if(res < _max) _max = res , rt = x;
}
map <int , int> M[2];
map <int , int> C;
struct node {
	int val , ty;
}s[MAXN];
int tot;
void dfs(int x , int fa , int val) {
	s[++tot].val = val;
	if(C[val]) s[tot].ty = 1; //C有值则它包含有平衡路径
	else s[tot].ty = 0;
	if(val == 0 && C[val] > 1) ans ++;//特判与根节点合法的路径
	C[val] ++; 
	for (int i = head[x]; i; i = nxt[i]) {
		if(to[i] == fa || vis[to[i]]) continue;
		dfs(to[i] , x , val + edge[i]);
	}
	C[val] --;
}
void work(int x) {
	get_siz(x , 0);
	rt = 0 , _max = 1e9;
	get_root(x , 0 , siz[x]);
	x = rt;vis[x] = 1;
	C[0] = 1; 
	M[0].clear();M[1].clear();
	for (int i = head[x]; i; i = nxt[i]) {
		if(vis[to[i]]) continue;
		tot = 0;
		dfs(to[i] , x , edge[i]);
		for (int j = 1; j <= tot; ++j) {
			if(s[j].ty == 0) ans += M[1][-s[j].val];
        		//当不包含平衡路径时在以找到的路径中查找包含平衡路径且能与自己组成平衡路径。
			else ans += M[0][-s[j].val] + M[1][-s[j].val];
        		//当包含时，能组成平衡路径中包不包含都可以
		}
		for (int j = 1; j <= tot; ++j) M[s[j].ty][s[j].val] ++;
	}
	for (int i = head[x]; i; i = nxt[i]) {
		if(vis[to[i]]) continue;
		work(to[i]);
	}
}
int main() {
	read(n);
	for (int i = 1; i < n; ++i) {
		int u , v , w;
		read(u);read(v);read(w);
		if(w == 0) w = -1;//这样做后，当一条路径是平衡路径则它的权值和为0
		add(u , v , w);add(v , u , w);
	}
	work(1);
	printf("%lld" , ans);
	return 0;
}