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