[BZOJ3697]采药人的路径（点分治）

=== ===

这里放传送门

=== ===

题解

首先考虑如果没有那个休息站的限制的话应该怎么做，可以想到把边权为0的变成边权为-1，这样的话就是找有多少条边权和为0的路径，这个东西非常好办。但是如果加入休息站的话要求这段路径的某一段前缀边权也为0。

考虑如果存在休息站的话这条路径有什么特点，只考虑经过当前子树根节点的路径，如果休息站恰好在根节点上，那么这条路径一定是两条deep为0的路径拼起来的；否则在这条路径的左侧或者右侧一定存在一个点，它的deep值和端点的deep值相等。

那么可以在对子树dfs的时候开一个数组表示这个deep值出现过没有。如果一条路径上已经出现了相等的deep，那么它可以跟任意路径拼成合法路径；但是如果是没有出现过相等deep的路径，那么只能是两条权值和为0的路径拼在一起。所以两种路径要分开存储分别计算。

代码

#include<cstdio>
#include<cstring>
#include<algorithm>
#define base 100000
#define inf 1000000000
using namespace std;
int n,p[100010],a[200010],nxt[200010],w[200010],deep[100010],size[100010],d[100010],c2;
int son[100010],ext[200010],root,cnt,tot;
long long ans;
bool vis[100010];
void add(int x,int y,int val){
tot++;a[tot]=y;nxt[tot]=p[x];w[tot]=val;p[x]=tot;
}
void getroot(int u,int fa,int N){
size[u]=1;son[u]=0;
for (int i=p[u];i!=0;i=nxt[i])
if (a[i]!=fa&&vis[a[i]]==false){
getroot(a[i],u,N);
size[u]+=size[a[i]];
son[u]=max(son[u],size[a[i]]);
}
son[u]=max(son[u],N-size[u]);
if (son[u]<son[root]) root=u;
}
void getdeep(int u,int fa,int dep,int dlt){
if (ext[dep+base]!=0)
deep[++cnt]=dep;
else if (ext[base]!=0) d[++c2]=dep;//根节点的空路径不能加进去
if (dep==0&&ext[dep+base]>=2) ans+=(long long)dlt;
ext[dep+base]++;
for (int i=p[u];i!=0;i=nxt[i])
if (a[i]!=fa&&vis[a[i]]==false)
getdeep(a[i],u,dep+w[i],dlt);
ext[dep+base]--;
}
long long calc(int u,int d0,int dlt){
long long sum=0;
cnt=c2=0;getdeep(u,0,d0,dlt);
sort(deep+1,deep+cnt+1);
sort(d+1,d+c2+1);
deep[0]=d[0]=-inf;
for (int l=1,r=cnt,p=r;l<r;){
while (p>l&&deep[p]>=deep[r]) --p;
p=max(p,l);//用一个指针找到权值相等的部分，注意p不能小于l。
if (deep[l]+deep[r]==0){
sum+=(long long)(r-p);++l;
}else
if (deep[l]+deep[r]>0) --r;
else ++l;
}//计算存在相等前缀的路径
for (int l=1,r=cnt,p=r;l<=c2&&r>=1;){
while (p>=1&&deep[p]>=deep[r]) --p;
if (d[l]+deep[r]==0){
sum+=(long long)(r-p);++l;
}else
if (d[l]+deep[r]>0) --r;
else ++l;
}//计算没有相等前缀的路径
return sum;
}
void work(int u){
ans+=calc(u,0,1);
vis[u]=true;
for (int i=p[u];i!=0;i=nxt[i])
if (vis[a[i]]==false){
ext[base]++;
ans-=calc(a[i],w[i],-1);
ext[base]--;
root=0;getroot(a[i],u,size[a[i]]);
work(root);
}
}
int main()
{
scanf("%d",&n);
for (int i=1;i<n;i++){
int x,y,v;
scanf("%d%d%d",&x,&y,&v);
v=(v==0)?-1:1;
add(x,y,v);add(y,x,v);
}
root=0;son[0]=inf;
getroot(1,0,n);
work(root);
printf("%I64d\n",ans);
return 0;
}