[BZOJ3697]采药人的路径

Posted by Dispwnl on February 20, 2019

题目

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。

题解

记阴性路径为$-1$,阳性为$1$,对于每个要处理的重心,统计子树中到根(重心)前缀和为$x$的点数量,$x\in[-n,n]$

设$f_{0,x}$为当前子树中前缀和为$x$且没有$x$的点数量,$f_{1,x}$为当前子树中前缀和为$x$且前面$x$出现过的点数量

$g​$表示其他子树的信息

那么这个重心的答案即为$f_{0,0}\times g_{0,0}+\sum_{i=-d}^{d}f_{0,i}\times g_{1,-i}+f_{1,i}\times (g_{0,-i}+g_{1,-i})$

发现似乎有一种情况漏求了?

重心开始的一条链上,$x,y$到重心的前缀和为$0$,这样的情况也应该计入答案

如果一个点符合条件答案就$+1$

注意两条路径不相同只跟路径端点有关,跟休息站位置无关(……)

代码

# include<iostream>
# include<cstring>
# include<cstdlib>
# include<cstdio>
# include<algorithm>
# define LL long long
using namespace std;
const int MAX=2e5+5;
struct p{
	int x,y;
	bool fl;
}c[MAX];
int n,num,rt,Sum,maxn;
LL ans;
int h[MAX],siz[MAX],Num[MAX],f[MAX];
int F[2][MAX],G[2][MAX];
bool use[MAX];
void add(int x,int y,bool fl)
{
	c[++num]=(p){h[x],y,fl},h[x]=num;
	c[++num]=(p){h[y],x,fl},h[y]=num;
}
int max(int x,int y) {return x>y?x:y;}
void GET_ROOT(int x=1,int fa=0)
{
	siz[x]=1,f[x]=0;
	for(int i=h[x];i;i=c[i].x)
	  if(!use[c[i].y]&&(c[i].y^fa)) GET_ROOT(c[i].y,x),siz[x]+=siz[c[i].y],f[x]=max(f[x],siz[c[i].y]);
	f[x]=max(f[x],Sum-siz[x]);
	if(f[x]<f[rt]) rt=x;
}
void Dfs(int x,int fa,int D1=0,int D2=0)
{
	maxn=max(maxn,D2+D1+1);
	if(Num[D2-D1+n])
	{
		++F[1][D2-D1+n];
		if(D2==D1) ++ans;
	}
	else ++F[0][D2-D1+n];
	++Num[D2-D1+n];
	for(int i=h[x];i;i=c[i].x)
	  if((c[i].y^fa)&&!use[c[i].y]) Dfs(c[i].y,x,D1+(c[i].fl^1),D2+c[i].fl);
	--Num[D2-D1+n];
}
void ask(int x)
{
	int D=0;
	for(int i=h[x];i;i=c[i].x)
	  if(!use[c[i].y])
	  {
	  	maxn=0,Dfs(c[i].y,x,c[i].fl^1,c[i].fl),D=max(D,maxn),ans+=1ll*G[0][n]*F[0][n];
	  	for(int i=-maxn;i<=maxn;++i)
	  	  ans+=1ll*F[0][n+i]*G[1][n-i]+1ll*F[1][n+i]*G[0][n-i]+1ll*F[1][n+i]*G[1][n-i],G[0][n-i]+=F[0][n-i],G[1][n-i]+=F[1][n-i];
	  	for(int i=-maxn;i<=maxn;++i)
	  	  F[0][n+i]=F[1][n+i]=0;
	  }
	for(int i=-D;i<=D;++i)
	  G[0][n+i]=G[1][n+i]=0;
}
void dfs(int x=rt,int fa=0)
{
	use[x]=1,ask(x);
	int xsum=Sum;
	for(int i=h[x];i;i=c[i].x)
	  if(!use[c[i].y]&&(c[i].y^fa))
	  {
	  	if(siz[c[i].y]>siz[x]) Sum=xsum-siz[x];
	  	else Sum=siz[c[i].y];
	  	f[rt=0]=Sum,GET_ROOT(c[i].y,x),dfs(rt,x);
	  }
}
int read()
{
	int x(0);
	char ch=getchar();
	for(;!isdigit(ch);ch=getchar());
	for(;isdigit(ch);x=x*10+ch-48,ch=getchar());
	return x;
}
int main()
{
	n=read();
	for(int i=1,x,y;i<n;++i)
	  x=read(),y=read(),add(x,y,read());
	return f[rt]=Sum=n,GET_ROOT(),dfs(),printf("%lld\n",ans),0;
}