[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;
}