BZOJ3697: 采药人的路径

Description

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

Input

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

Output

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

Sample Input

7
1 2 0
3 1 1
2 4 0
5 2 0
6 3 1
5 7 1

Sample Output

1

HINT

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

Source

点分治 把边权0变为-1

记录f[i][0/1]表示当前子树中dis为i的有多少个 0/1表示是否是第一次出现（即分割点）

g[i][0/1]表示之前已经考虑过的子树

然后就没了

debug了一个小时最后发现我把g[i][0/1]每次结束后置为了1？****

#include<bits/stdc++.h>

using namespace std;

const int maxn=100010;

int dis[maxn],size[maxn],dep[maxn];

int t[maxn<<1],n,mxdep,head[maxn],cnt,fa[maxn];

int f[maxn<<1][2],g[maxn<<1][2];

bool vis[maxn];

long long ans;

struct edge
{
	int to,nxt,val;
}e[maxn<<1];

inline void addedge(int x,int y,int w)
{
	e[++cnt].to=y;
	e[cnt].nxt=head[x];
	head[x]=cnt;
	e[cnt].val=w?1:-1;
}

void getroot(int x,int Size,int &cg)
{
	bool flag=true;
	size[x]=1;
	for(int i=head[x];i;i=e[i].nxt)
	{
		int y=e[i].to;
		if(!vis[y]&&fa[x]!=y)
		{
			fa[y]=x;
			getroot(y,Size,cg);
			size[x]+=size[y];
			if(size[y]>Size/2) flag=false;
		}
	}
	if(size[x]<Size/2) flag=false;
	if(flag) cg=x;
}

void dfs(int x)
{
	size[x]=1;
	mxdep=max(mxdep,dep[x]);
	if(t[dis[x]]) f[dis[x]][1]++;
	else f[dis[x]][0]++;
	t[dis[x]]++;
	for(int i=head[x];i;i=e[i].nxt)
	{
		int y=e[i].to;
		if(!vis[y]&&fa[x]!=y)
		{
			fa[y]=x;
			dis[y]=dis[x]+e[i].val;
			dep[y]=dep[x]+1;
			dfs(y);
			size[x]+=size[y];
		}
	}
	t[dis[x]]--;
}

void solve(int x,int Size)
{
	int cg,mx=0;
	fa[x]=0;
	getroot(x,Size,cg);
	vis[cg]=1;size[fa[cg]]=Size-size[cg];
	fa[cg]=0;dis[cg]=n;g[n][0]=1;
	for(int i=head[cg];i;i=e[i].nxt)
		if(!vis[e[i].to])
		{
			dis[e[i].to]=dis[cg]+e[i].val;
			dep[e[i].to]=1;
			fa[e[i].to]=cg;
			mxdep=1;dfs(e[i].to);mx=max(mx,mxdep);
			ans+=1ll*(g[n][0]-1)*f[n][0];
			for(int j=-mxdep;j<=mxdep;j++)
				ans+=1ll*g[n-j][0]*f[n+j][1]+1ll*g[n-j][1]*f[n+j][0]+1ll*g[n-j][1]*f[n+j][1];
			for(int j=n-mxdep;j<=n+mxdep;j++)
				g[j][0]+=f[j][0],f[j][0]=0,g[j][1]+=f[j][1],f[j][1]=0;
		}
	for(int i=n-mx;i<=n+mx;i++) g[i][0]=g[i][1]=0;
	for(int i=head[cg];i;i=e[i].nxt)
		if(!vis[e[i].to])
			solve(e[i].to,size[e[i].to]);
}

int main()
{
	scanf("%d",&n);
	for(int i=1;i<n;i++) { int x,y,w; scanf("%d%d%d",&x,&y,&w); addedge(x,y,w); addedge(y,x,w); }
	solve(1,n);
	return printf("%lld\n",ans),0;
}