[JZOJ3234] 阴阳

题目

题目大意

有一棵树，每条边有$0$或$1$两种颜色。
求满足存在$(u,v)$路径上的点$x$使得$(u,x)$和$(x,v)$路径上的两种颜色出现次数相同的点对$(u,v)$数量。

---

思考历程

在看到这题之前我就已经知道这题是点分治了……
然而看了题目后，很长一段时间搞错了题目大意……
后来终于搞懂了题目大意，于是开始想正解。
想了半天，心态崩了……点分治我在差不多一年前才打过一题啊！
很久后搜题解，粗略看一看，没有一个看得懂……
于是又开始自己刚……然后刚出来了……
结果WA了……
调了不知道多久，终于AC……

---

正解

正解就是点分治啦……
显然的思路是将$0$颜色的边权值记为$1$，$1$颜色的边权值记为$-1$。
如果$(u,v)$满足条件，一定有点$x$使得$(u,x)$和$(x,v)$路径上的权值和为$0$。
现在设重心为$root$，现在我们要求的路径必须经过$root$。
求出每个点到$root$路径上的权值和$di{s}_x$。
显然，如果$(u,v)$满足条件，一定有$di{s}_u+di{s}_v=0$。
还有一个条件是存在$u$或$v$的祖先$x$，使得$di{s}_x=di{s}_u$或$di{s}_x=di{s}_v$。
然后我们就搬出几个桶……
设$to{t}_i$表示$dis$值为$i$的点的数量，$ca{n}_i$表示$dis$值为$i$并且其祖先有$dis$值和它相同的点的数量。这两个都是前面计算过的$root$的所有子树的。同理，设$su{b}_i$和$sc{n}_i$，意义相同，表示现在正在计算的这个子树的。
$su{b}_i$的计算方法显然，至于$sc{n}_i$，我们在$dfs$的过程中再维护一个桶$an{c}_i$，表示$dis$值为$i$的祖先有多少个，这样就可以快速判断了。
$to{t}_i$和$ca{n}_i$可以由后两者累加而得。
计算答案的时候，枚举$dis$值，用前面四个桶里的值来计算，具体来说，如果$(u,v)$能被算入答案中，则$u$和$v$至少有一个在$can$或$scn$中。
注意，我们还没有计算$v=root$的情况。针对这种情况，我们维护一个$cnt0$，在$dfs$的时候遇到$an{c}_{di{s}_i}>1$时，便意味着除了$root$之外，还有一个和它$dis$值相同的祖先，那就将其加入$cnt0$中。答案加上$cnt0$即可。

所以说实际上这是一道点分治的模板题啊……

---

代码

using namespace std;
#include <cstdio>
#include <cstring>
#include <algorithm>
#include <climits>
#include <cassert>
#define N 100010
int n;
struct EDGE{
	int to,len;
	EDGE *las;
	bool ok;
} e[N*2];
int ne;
EDGE *last[N];
inline void link(int u,int v,int len){
	e[ne]={v,len,last[u],1};
	last[u]=e+ne++;
}
#define rev(ei) (e+(int((ei)-e)^1))
long long ans;
int siz[N],all;
void get_siz(int x,int fa){
	siz[x]=1;
	for (EDGE *ei=last[x];ei;ei=ei->las)
		if (ei->to!=fa && ei->ok){
			get_siz(ei->to,x);
			siz[x]+=siz[ei->to];
		}
}
int find(int x,int fa){//找重心
	bool bz=(all-siz[x]<<1<=all);
	for (EDGE *ei=last[x];ei;ei=ei->las)
		if (ei->to!=fa && ei->ok){
			int tmp=find(ei->to,x);
			if (tmp)
				return tmp;
			bz&=(siz[ei->to]<<1<=all);
		}
	if (bz)
		return x;
}
int dis[N];
int _tot[N*2],*tot=_tot+N;//为了代码方便，所以用指针来处理了……
int _sub[N*2],*sub=_sub+N;
int _can[N*2],*can=_can+N;
int _scn[N*2],*scn=_scn+N;
int cnt0;
int _anc[N*2],*anc=_anc+N;
int mn,mx,smn,smx;//记录mn,mx是为了方便清空桶。
void get_dis(int x,int fa){
	smn=min(smn,dis[x]),smx=max(smx,dis[x]);
	sub[dis[x]]++;
	if (anc[dis[x]]){
		scn[dis[x]]++;
		if (anc[dis[x]]>1 && dis[x]==0)
			cnt0++;
	}
	anc[dis[x]]++;
	for (EDGE *ei=last[x];ei;ei=ei->las)
		if (ei->to!=fa && ei->ok){
			dis[ei->to]=dis[x]+ei->len;
			get_dis(ei->to,x);
		}
	anc[dis[x]]--;
}
void divide(int x){
	get_siz(x,0);
 	all=siz[x];
	int root=find(x,0);
	dis[root]=0;
	mn=INT_MAX,mx=INT_MIN;
	anc[0]=1;
	for (EDGE *ei=last[root];ei;ei=ei->las)
		if (ei->ok){
			smn=INT_MAX,smx=INT_MIN;
			cnt0=0;
			dis[ei->to]=dis[root]+ei->len;
			get_dis(ei->to,root);
			for (int j=smn;j<=smx;++j)
				ans+=1ll*sub[j]*can[-j]+1ll*scn[j]*tot[-j]-1ll*scn[j]*can[-j];//容斥原理
			ans+=cnt0;
			for (int j=smn;j<=smx;++j){
				tot[j]+=sub[j],sub[j]=0;
				can[j]+=scn[j],scn[j]=0;
			}
			mn=min(mn,smn);
			mx=max(mx,smx);
		}
	anc[0]=0;
	for (int i=mn;i<=mx;++i)
		tot[i]=can[i]=0;
	for (EDGE *ei=last[root];ei;ei=ei->las)
		if (ei->ok){
			ei->ok=rev(ei)->ok=0;
			divide(ei->to);
		}
}
int main(){
	scanf("%d",&n);
	for (int i=1;i<n;++i){
		int u,v,type;
		scanf("%d%d%d",&u,&v,&type);
		link(u,v,type==0?1:-1);
		link(v,u,type==0?1:-1);
	}
	divide(1);
	printf("%lld\n",ans);
	return 0;
}

---

总结

点分治的题目还是打得太少了……