2019.01.09 bzoj3697: 采药人的路径（点分治）

传送门
点分治好题。
题意：给出一棵树，边分两种，求满足由两条两种边数相等的路径拼成的路径数。

---

思路：
考虑将边的种类转化成边权$-1$−1和$1$1，这样就只用考虑由两条权值为$0$0的路径拼成的路径数。
然后发现对于一个点$p$p，经过它的这样的路径满足如下至少一个条件：

1. $\exist v$∃v满足$dist(v,p)=0\&\&\exist v_0\in road_{v,p}$dist(v,p)=0&&∃v0​∈roadv,p​满足$dist(v_0,v)=0$dist(v0​,v)=0
2. $\exist v_1,v_2$∃v1​,v2​满足$dist(v_1,p)=0,dist(v_2,p)=0
3. $\exist v_1,v_2$∃v1​,v2​满足$dist(v_1,p)+dist(v_2,p)=0\&\&\exist v_0\in road_{v_1,p}$dist(v1​,p)+dist(v2​,p)=0&&∃v0​∈roadv1​,p​满足$dist{v_0,p}=dist_{v_1,p}$distv0​,p=distv1​,p​
4. $\exist v_1,v_2$∃v1​,v2​满足$dist(v_1,p)+dist(v_2,p)=0\&\&\exist v_0\in road_{v_2,p}$dist(v1​,p)+dist(v2​,p)=0&&∃v0​∈roadv2​,p​满足$dist{v_0,p}=dist_{v_2,p}$distv0​,p=distv2​,p​

然后点分容斥处理即可，注意不要忘记第一种情况的讨论。
代码：

#include<bits/stdc++.h>
#define fi first
#define se second
#define ri register int
using namespace std;
const int N=1e5+5;
int n,siz[N],msiz[N],rt=0,all=0,lim;
bool vis[N];
typedef pair<int,int> pii;
typedef long long ll;
vector<pii>e[N];
ll ans=0,cnt[2][N][2]={0};
inline int read(){
	int ans=0;
	char ch=getchar();
	while(!isdigit(ch))ch=getchar();
	while(isdigit(ch))ans=(ans<<3)+(ans<<1)+(ch^48),ch=getchar();
	return ans;
}
void findroot(int p,int fa){
	siz[p]=1,msiz[p]=0;
	for(ri i=0,v;i<e[p].size();++i){
		if((v=e[p][i].fi)==fa||vis[v])continue;
		findroot(v,p),siz[p]+=siz[v],msiz[p]=max(msiz[p],siz[v]);
	}
	msiz[p]=max(msiz[p],all-siz[p]);
	if(msiz[p]<msiz[rt])rt=p;
}
void getdis(int p,int fa,int dist,int down,int up){
	if(dist>=down&&dist<=up){
		if(dist>=0)++cnt[0][dist][1];
		if(dist<0)++cnt[1][-dist][1];
	}
	else{
		if(dist>=0)++cnt[0][dist][0];
		if(dist<0)++cnt[1][-dist][0];
	}
	down=min(down,dist),up=max(up,dist),lim=max(lim,max(-down,up));
	for(ri i=0,v;i<e[p].size();++i){
		if((v=e[p][i].fi)==fa||vis[v])continue;
		getdis(v,p,dist+e[p][i].se,down,up);
	}
}
inline void calc(int p,int pre,int tmp){
	ll sum=0;
	lim=0; 
	if(tmp==1)getdis(p,0,pre,1,-1);
	else getdis(p,0,pre,0,0);
	sum+=cnt[0][0][1]*(cnt[0][0][1]-1)/2,cnt[0][0][0]=cnt[0][0][1]=0;
	for(ri i=1;i<=lim;++i){
		sum+=cnt[0][i][0]*cnt[1][i][1];
		sum+=cnt[0][i][1]*cnt[1][i][0];
		sum+=cnt[0][i][1]*cnt[1][i][1];
		cnt[0][i][0]=cnt[0][i][1]=cnt[1][i][0]=cnt[1][i][1]=0;
	}
	ans+=sum*tmp;
}
void solve(int p,int fa,int tim,int dist){
	if(!dist){
		if(tim>=2)++ans;
		++tim;
	}
	for(ri i=0,v;i<e[p].size();++i)if((v=e[p][i].fi)!=fa&&!vis[v])solve(v,p,tim,dist+e[p][i].se);
}
void dfs(int p){
	vis[p]=1,solve(p,0,0,0),calc(p,0,1);
	for(ri i=0,v;i<e[p].size();++i){
		if(vis[v=e[p][i].fi])continue;
		calc(v,e[p][i].se,-1),rt=0,all=siz[v],findroot(v,0),dfs(rt);
	}
}
int main(){
	freopen("lx.in","r",stdin);
	n=read();
	for(ri i=1,u,v,w;i<n;++i)u=read(),v=read(),w=read()?1:-1,e[u].push_back(pii(v,w)),e[v].push_back(pii(u,w));
	rt=0,all=n,msiz[rt]=n,findroot(1,0),dfs(rt);
	cout<<ans;
	return 0;
}