bzoj 3697: 采药人的路径 点分治

题意

采药人的药田是一个树状结构，每条路径上都种植着同种药材。

采药人以自己对药材独到的见解，对每种药材进行了分类。大致分为两类，一种是阴性的，一种是阳性的。

采药人每天都要进行采药活动。他选择的路径是很有讲究的，他认为阴阳平衡是很重要的，所以他走的一定是两种药材数目相等的路径。采药工作是很辛苦的，所以他希望他选出的路径中有一个可以作为休息站的节点（不包括起点和终点），满足起点到休息站和休息站到终点的路径也是阴阳平衡的。他想知道他一共可以选择多少种不同的路径。

N ≤ 100,000。

分析

一开始在点分的时候求的是以每个分治中心为休息站的路径数，愉快地码完后发现这样是错的。我们要求的是经过一个分治中心的满足条件的路径数。这样的话就复杂很多了。

我们在递归的时候，需要维护s1[x][0/1],s2[x][0/1]，0表示负数，1表示非负数，s1[x][0/1]表示路径权值为x且路上有休息站的路径数，s2则表示没有休息站的路径数。怎么判断路上是否有休息站呢？根据定义，那么只要从终点到路径上某个点的路径长度为0，则有休息站。这时候我们可以维护一下dfs栈内的信息，若祖先中有一个点的权值与其相同，则表示其到该祖先的路径长度为0，也就是有休息站。然后每次把路径信息合并一下即可。

细节有点小多。

代码

#include<iostream>
#include<cstdio>
#include<cstdlib>
#include<cstring>
#include<algorithm>
using namespace std;

typedef long long LL;

const int N=100005;

int n,last[N],cnt,size[N],val[N],tot,root,f[N][2],g[N][2],s1[N][2],s2[N][2],a[N],now,ins[N][2];
struct edge{int to,next,w;}e[N*2];
bool vis[N],use[N][2];
LL ans;

int read()
{
int x=0,f=1;char ch=getchar();
while (ch<'0'||ch>'9'){if(ch=='-')f=-1;ch=getchar();}
while (ch>='0'&&ch<='9'){x=x*10+ch-'0';ch=getchar();}
return x*f;
}

void addedge(int u,int v,int w)
{
e[++cnt].to=v;e[cnt].w=w;e[cnt].next=last[u];last[u]=cnt;
e[++cnt].to=u;e[cnt].w=w;e[cnt].next=last[v];last[v]=cnt;
}

void get_root(int x,int fa)
{
size[x]=1;val[x]=0;
for (int i=last[x];i;i=e[i].next)
{
if (e[i].to==fa||vis[e[i].to]) continue;
get_root(e[i].to,x);
size[x]+=size[e[i].to];
val[x]=max(val[x],size[e[i].to]);
}
val[x]=max(val[x],tot-size[x]);
if (!root||val[x]<val[root]) root=x;
}

void get_len(int x,int fa,int len)
{
a[++now]=len;
size[x]=1;
if (ins[abs(len)][len>=0]) s1[abs(len)][len>=0]++;
else s2[abs(len)][len>=0]++;
ins[abs(len)][len>=0]++;
for (int i=last[x];i;i=e[i].next)
if (e[i].to!=fa&&!vis[e[i].to]) get_len(e[i].to,x,len+e[i].w),size[x]+=size[e[i].to];
ins[abs(len)][len>=0]--;
}

void solve(int x)
{
vis[x]=1;
for (int i=last[x];i;i=e[i].next)
{
if (vis[e[i].to]) continue;
int tmp=now;
get_len(e[i].to,x,e[i].w);
for (int j=tmp+1;j<=now;j++)
{
int x=a[j];
if (use[abs(x)][x>=0]) continue;
ans+=(LL)s1[abs(x)][x>=0]*(f[abs(x)][x<=0]+g[abs(x)][x<=0]);
ans+=(LL)s2[abs(x)][x>=0]*f[abs(x)][x<=0];
if (!x) ans+=(LL)s2[0][1]*g[0][1];
use[abs(x)][x>=0]=1;
}
for (int j=tmp+1;j<=now;j++)
{
int x=a[j];
use[abs(x)][x>=0]=0;
f[abs(x)][x>=0]+=s1[abs(x)][x>=0];
g[abs(x)][x>=0]+=s2[abs(x)][x>=0];
s1[abs(x)][x>=0]=s2[abs(x)][x>=0]=0;
}
}
ans+=f[0][1];
while (now)
{
int x=a[now];now--;
f[abs(x)][x>=0]=g[abs(x)][x>=0]=0;
}
for (int i=last[x];i;i=e[i].next)
{
if (vis[e[i].to]) continue;
root=0;tot=size[e[i].to];
get_root(e[i].to,x);
solve(root);
}
}

int main()
{
n=read();
for (int i=1;i<n;i++)
{
int x=read(),y=read(),z=read();
if (!z) z=-1;
addedge(x,y,z);
}
tot=n;root=0;
get_root(1,0);
solve(root);
printf("%lld",ans);
return 0;
}