bzoj3697 采药人的路径（点分治+dp）

把边权0变成-1，路径边权和为0就是阴阳平衡了。我们点分治，考虑一条过根的路径，被根分成两部分，如果是合法的，那么休息站不是在前半段，就是在后半段。我们记g[i][0]表示到根的边权和为i，且不存在休息站的路径数；g[i][1]表示存在休息站的路径数。一个一个子树的做，避免不合法的情况。用tmp[i][0/1]表示已经做过的子树的信息。则对答案的贡献就是
g[0][0]∗tmp[0][0]+∑i=−mxdepmxdepg[i][1]∗tmp[−i][0]+g[i][0]∗tmp[−i][1]+g[i][1]∗tmp[−i][1]。注意特殊处理一下根的情况。

#include <cstdio>
#include <cstring>
#include <algorithm>
using namespace std;
#define ll long long
#define inf 0x3f3f3f3f
#define N 100010
inline 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;
}
int n,h[N],num=0,rt,sz[N],sumsz,f[N],fa[N],tmp[N<<1][2],dis[N],g[N<<1][2],cnt[N<<1],dep[N],mxdep;
bool vis[N];ll ans=0;//g[i][0],表示到根的边权和为i，且不存在休息站的路径数；g[i][1]表示存在休息站的路径数
struct node{
    int to,next,val;
}data[N<<1];
void dfs1(int x){
    sz[x]=1;
    for(int i=h[x];i;i=data[i].next){
        int y=data[i].to;if(y==fa[x]||vis[y]) continue;
        fa[y]=x;dep[y]=dep[x]+1;dfs1(y);sz[x]+=sz[y];
    }
}
void dfs2(int x){
    f[x]=0;
    for(int i=h[x];i;i=data[i].next){
        int y=data[i].to;if(y==fa[x]||vis[y]) continue;
        dfs2(y);f[x]=max(f[x],sz[y]);
    }f[x]=max(f[x],sumsz-sz[x]);if(f[x]<f[rt]) rt=x;
}
void dfs3(int x){
    int val=dis[x];mxdep=max(mxdep,dep[x]);
    if(cnt[val]) g[val][1]++;
    else g[val][0]++;
    cnt[val]++;
    for(int i=h[x];i;i=data[i].next){
        int y=data[i].to;if(y==fa[x]||vis[y]) continue;
        dis[y]=dis[x]+data[i].val;dfs3(y);
    }cnt[val]--;
}
inline void calc(int x){
    tmp[n][0]=1;int mx=0;
    for(int ii=h[x];ii;ii=data[ii].next){
        int y=data[ii].to;if(vis[y]) continue;
        mxdep=0;dis[y]=data[ii].val+n;dfs3(y);mx=max(mx,mxdep);
        ans+=(ll)g[n][0]*(tmp[n][0]-1);//以x为休息站
        for(int i=-mxdep;i<=mxdep;++i)
            ans+=(ll)g[n+i][0]*tmp[n-i][1]+(ll)g[n+i][1]*tmp[n-i][0]+(ll)g[n+i][1]*tmp[n-i][1];
        for(int i=n-mxdep;i<=n+mxdep;++i)
            for(int j=0;j<2;++j)
                tmp[i][j]+=g[i][j],g[i][j]=0;
    }for(int i=n-mx;i<=n+mx;++i) tmp[i][0]=tmp[i][1]=0;
}
inline void solve(int x){
    vis[x]=1;fa[x]=0;dep[x]=0;dfs1(x);calc(x);
    for(int i=h[x];i;i=data[i].next){
        int y=data[i].to;if(vis[y]) continue;
        rt=0;sumsz=sz[y];dfs2(y);solve(rt);
    }
}
int main(){
//  freopen("a.in","r",stdin);
    n=read();f[0]=inf;
    for(int i=1;i<n;++i){
        int x=read(),y=read(),val=read();if(!val) val=-1;
        data[++num].to=y;data[num].next=h[x];h[x]=num;data[num].val=val;
        data[++num].to=x;data[num].next=h[y];h[y]=num;data[num].val=val;
    }dfs1(1);rt=0;sumsz=n;dfs2(1);solve(rt);
    printf("%lld\n",ans);
    return 0;
}