BZOJ 3697 采药人的路径 树的点分治

题目大意：给出一棵n个结点的树，边权为1或-1，求有多少条路径边权和为0并且能在路径上找到一个在起点与终点之间的休息点(不能是起点或终点)满足起点到休息点、休息点到终点的路径边权为0

不是那么裸的点分治。

首先考虑在一棵子树中如何找到边权和为0的路径。由于边权只有1和-1，和的绝对值不会超过n，记录一下权值为i的路径分别有多少条，其中i属于[-n,n]。这样i与-i组合形成的路径边权和为0。

再考虑找休息点。在子树中由根拼接两条路径组成的一整条路径中，休息点可能在三个位置：起点所在的半条路径，终点所在的半条路径，根。

这样的话记录两个数组：f(i,j)表示之前扫过的子树中权值为i的路径的条数,j代表在根与终点之间是否有某个点满足从根到这个点的权值为i；g(i,j)表示当前的子树，其余与f的描述相同。

则统计答案时枚举i，分别按照休息点在起点所在的半条路径( g(i,1) * f(-i,0) )，休息点在终点所在的半条路径( g(i,0) * f(-i,1) ) ，休息点在两端均可( g(i,1) * f(-i,1) )，休息点在根( (g(0,0)-1) * f(0,0) )统计答案即可。

#include <cstdio>
#include <cstring>
#include <algorithm>
#define N 100005
using namespace std;
typedef long long LL;
struct Edge {
    int from,to,nxt,val;
    Edge() {}
    Edge(int _from,int _to,int _nxt,int _val):
        from(_from),to(_to),nxt(_nxt),val(_val) {}
}e[N*2];
int n,tot=-1,fir[N];
void Add_Edge(int x,int y,int v) {
    e[++tot]=Edge(x,y,fir[x],v), fir[x]=tot;
    e[++tot]=Edge(y,x,fir[y],v), fir[y]=tot;
    return ;
}
int sum,root,max_dpt,dpt[N],dist[N],pa[N],max_siz[N],siz[N],t[N*2];
bool vis[N];
LL ans,f[N*2][2],g[N*2][2];
void dfs_root(int x,int from) {
    pa[x]=from, siz[x]=1, max_siz[x]=0;
    for(int i=fir[x];~i;i=e[i].nxt) {
        if(e[i].to==from || vis[e[i].to]) continue;
        dfs_root(e[i].to,x);
        siz[x]+=siz[e[i].to];
        max_siz[x]=max(max_siz[x],siz[e[i].to]);
    }
    max_siz[x]=max(max_siz[x],sum-siz[x]);
    if(max_siz[x]<max_siz[root]) root=x;
    return ;
}
void get_root(int x) {
    root=0;
    sum=siz[x];
    dfs_root(x,0);
    siz[pa[root]]=sum-siz[root];
    return ;
}
void dfs_dpt(int x,int from) {
    max_dpt=max(max_dpt,dpt[x]);
    if(t[dist[x]]) f[dist[x]][1]++;
    else f[dist[x]][0]++;
    t[dist[x]]++;
    for(int i=fir[x];~i;i=e[i].nxt) {
        if(e[i].to==from || vis[e[i].to]) continue;
        dpt[e[i].to]=dpt[x]+1;
        dist[e[i].to]=dist[x]+e[i].val;
        dfs_dpt(e[i].to,x);
    }
    t[dist[x]]--;
    return ;
}
void calc(int x) {
    int maxx=0;
    g[n][0]=1;
    for(int i=fir[x];~i;i=e[i].nxt) {
        if(vis[e[i].to]) continue;
        dist[e[i].to]=n+e[i].val;
        dpt[e[i].to]=1;
        max_dpt=1;
        dfs_dpt(e[i].to,0);
        maxx=max(max_dpt,maxx);
        ans+=(g[n][0]-1)*f[n][0];
        for(int j=-max_dpt;j<=max_dpt;j++)
            ans+=g[n-j][1]*f[n+j][1]+g[n-j][0]*f[n+j][1]+g[n-j][1]*f[n+j][0];
        for(int j=n-max_dpt;j<=n+max_dpt;j++)
            g[j][0]+=f[j][0], g[j][1]+=f[j][1], f[j][0]=f[j][1]=0;
    }
    for(int i=n-maxx;i<=n+maxx;i++) g[i][0]=g[i][1]=0;
    return ;
}
void Divide(int x) {
    vis[x]=true;
    calc(x);
    for(int i=fir[x];~i;i=e[i].nxt) {
        if(vis[e[i].to]) continue;
        get_root(e[i].to);
        Divide(root);
    }
    return ;
}
int main() {
    memset(fir,-1,sizeof fir);
    scanf("%d",&n);
    for(int i=1;i<n;i++) {
        int x,y,z;
        scanf("%d%d%d",&x,&y,&z);
        if(!z) z=-1;
        Add_Edge(x,y,z);
    }
    max_siz[0]=siz[1]=n;
    get_root(1);
    Divide(root);
    printf("%lld\n",ans);
    return 0;
}