题目
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;
}