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

自从gay豪讲了后缀自动机以后就不知道在干什么了，今天闲下来搞搞自己的专题，至于动态点分治什么的真的是坑啊。

好了，回到这道题，同样是处理路径问题，想到树分治。首先考虑没有休息站的方法，这里我想到有两种，一种就是dfs是处理出dis数组，排序后用数学神马的随便搞搞，还有一种就是基于dp思想的，定义数组f[i],g[i]分别表示当前子树中路径长度为i的节点数和之前子树中路径为i的节点数，这样就可以保证求出来的路径一定是经过根节点的。但是这道题除了这个限制以外还有休息站，多一维表示这条路径中有没有休息站，因为休息站一定在两条链之间，在dfs的时候就记录一下当前路径的前缀有那些路径和时出现过就好了。

    #include<cstdio>  
    #include<cstring>  
    #include<iostream>  
    #define LL long long  
    #include<cmath>  
    #include<cstdlib>  
    #define maxn 200020  
    using namespace std;  
    int head[maxn],tot,n,ff[maxn],sz[maxn],rt,size,vis[maxn],cnt[maxn];  
    LL f[maxn][2],g[maxn][2],ans;  
    //f[i][0]表示当前子树中和为i且没有休息站的 方案数1是有休息站  
    struct edge{int v,next,w;}e[maxn];   
    void adde(int a,int b,int c){e[tot].v=b,e[tot].next=head[a],e[tot].w=c;head[a]=tot++;}  
      
    void getrt(int u,int fa){  
        ff[u]=0,sz[u]=1;  
        for(int v,i=head[u];i!=-1;i=e[i].next){  
            if(vis[v=e[i].v]||v==fa)continue;  
            getrt(v,u);  
            sz[u]+=sz[v];  
            ff[u]=max(ff[u],sz[v]);  
        }  
        ff[u]=max(ff[u],size-ff[u]);  
        if(ff[u]<ff[rt])rt=u;  
    }  
    int Max;  
    void dfs(int u,int fa,int sum){  
        if(cnt[sum])f[sum][0]++;  
        else f[sum][1]++;  
        cnt[sum]++;  
        int x=abs(sum-n);  
        Max=max(x,Max);  
        for(int v,i=head[u];i!=-1;i=e[i].next){  
            if(vis[v=e[i].v]||v==fa)continue;  
            dfs(v,u,sum+e[i].w);  
        }  
        cnt[sum]--;  
    }  
      
    void calc(int u){  
        g[n][0]=1;int mx=-1e9;  
        for(int v,i=head[u];i!=-1;i=e[i].next){  
            if(vis[v=e[i].v])continue;Max=-n;  
            dfs(v,u,e[i].w+n);  
            mx=max(mx,Max);  
            ans+=(g[n][1]-1)*f[n][1];  
            for(int j=-Max;j<=Max;j++)  
                ans+=f[n-j][0]*g[n+j][0]+f[n-j][0]*g[n+j][1]+f[n-j][1]*g[n+j][0];  
            for(int j=n-Max;j<=n+Max;j++){  
                g[j][0]+=f[j][0],g[j][1]+=f[j][1];  
                f[j][0]=f[j][1]=0;  
            }  
        }  
        for(int i=n-mx;i<=n+mx;i++)g[i][0]=g[i][1]=0;  
    }  
      
    void solve(int u){  
        vis[u]=1;  
        calc(u);  
        for(int v,i=head[u];i!=-1;i=e[i].next){  
            if(vis[v=e[i].v])continue;  
            ff[rt=0]=size=sz[v];  
            getrt(v,u);  
            solve(rt);  
        }  
    }  
      
    int main(){  
        memset(head,-1,sizeof(head));  
        scanf("%d",&n);  
        for(int a,b,c,i=1;i<n;i++){  
            scanf("%d%d%d",&a,&b,&c);  
            if(c==1)adde(a,b,1),adde(b,a,1);  
            else adde(a,b,-1),adde(b,a,-1);  
        }  
        ff[rt=0]=size=n;  
        getrt(1,0);  
        solve(rt);  
        printf("%lld",ans);  
        return 0;  
    }