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

Description

采药人的药田是一个树状结构，每条路径上都种植着同种药材。

采药人以自己对药材独到的见解，对每种药材进行了分类。大致分为两类，一种是阴性的，一种是阳性的。

采药人每天都要进行采药活动。他选择的路径是很有讲究的，他认为阴阳平衡是很重要的，所以他走的一定是两种药材数目相等的路径。采药工作是很辛苦的，所以他希望他选出的路径中有一个可以作为休息站的节点（不包括起点和终点），满足起点到休息站和休息站到终点的路径也是阴阳平衡的。他想知道他一共可以选择多少种不同的路径。

Input

第1行包含一个整数N。

接下来N-1行，每行包含三个整数a_i、b_i和t_i，表示这条路上药材的类型。

Output

输出符合采药人要求的路径数目。

Sample Input

7

1 2 0

3 1 1

2 4 0

5 2 0

6 3 1

5 7 1

Sample Output

1

HINT

对于100%的数据，N ≤ 100,000。

Source

[ Submit][ Status][ Discuss] 

分析：
把阳性药的价值设为1，阴性药的价值为-1
如果没有所谓休息站的限制
我们只要计算路径和为0的路径条数即可（比较简单）

而现在我们需要在路径上找到一个休息站，怎么破？
我们求出每个结点到根节点的路径和$sum$，如果这条路径上有两个结点的$sum$相等
就说明有一段的路径和为0，即这个路径上可以设置一个休息站

$f[i][0/1]$表示当前子树中，子结点到根结点的路径长度为$i$，其中是否可以设置休息站
$g[i][0/1]$表示在已经处理过的子树中，子结点到根结点的路径长度为$i$，，其中是否可以设置休息站
（这样设置就可以不用去重了）

这个表述可能有点难理解（但是绝对正确）：

具体的计算方法：
如果在dfs一条路径时，
到达结点$now$之前有一个结点的$sum$值等于$sum[now]$
那么$now$就可以作为路径的一个端点，$f[sum[now]][1]++$

$ans={f[0][0]*g[0][0]+{\sum_{i=-deep}^{deep}}(f[i][0]*g[-i][1]+f[i][1]*g[-i][0]+f[i][1]*g[-i][1])}$

---

Refun学弟问了我一个非常高深的问题：
为什么在solve中，一开始：g[o][0]=1;
但是在统计答案的时候：ans+=(g[o][0]-1)*f[o][0];
这个g[o][0]=1;看似没有一点作用啊

实际上，这个初始化非常的重要
看代码中的转移方程：

ans=ans+f[j+o][1]*g[-j+o][0]+f[j+o][0]*g[-j+o][1]+f[j+o][1]*g[-j+o][1];

我们观察到第一项f[j+o][1]*g[-j+o]，当$j=0$的时候会用到$g[o][0]$
$f[j+o][1]$记录的是从根结点出发至少一步，子结点到根结点的路径长度为$i$，其中可以设置休息站的路径条数
那我们看一下当$j=0$时，f数组是个什么情况

显然这是一种合法情况（点2上设一个休息站即可）
用不到再找一条链拼成完整的一条路径，这样的路径也是经过当前根结点的合法路径
所以我们要把这一部分也累加上，因此$g[o][0]=1$（防止f[j+o][1]*g[-j+o]乘成0）

#include<bits/stdc++.h>
#define ll long long 

using namespace std;

const int N=100010;
int o,n,sz,root,F[N],size[N],st[N],tot=0;
ll f[N<<1][2],g[N<<1][2],ans=0;
int sum[N],deep[N];
struct node{
    int y,nxt,v;
};
node way[N<<1];
bool vis[N];

void add(int u,int w,int z) {
    tot++;way[tot].y=w;way[tot].v=z;way[tot].nxt=st[u];st[u]=tot;
    tot++;way[tot].y=u;way[tot].v=z;way[tot].nxt=st[w];st[w]=tot;
}

void findroot(int now,int fa) {
    size[now]=1;
    F[now]=0;
    for (int i=st[now];i;i=way[i].nxt)
        if (way[i].y!=fa&&!vis[way[i].y]) {
            findroot(way[i].y,now);
            size[now]+=size[way[i].y];
            F[now]=max(F[now],size[way[i].y]);
        }
    F[now]=max(F[now],sz-size[now]);
    if (F[now]<F[root]) root=now;
}

int mxdeep;
int cnt[N<<1];

void dfs(int now,int fa) {
    mxdeep=max(mxdeep,deep[now]);
    if (cnt[sum[now]+o]) f[sum[now]+o][1]++;
    else f[sum[now]+o][0]++;
    cnt[sum[now]+o]++;

    for (int i=st[now];i;i=way[i].nxt)
        if (way[i].y!=fa&&!vis[way[i].y]) {
            sum[way[i].y]=sum[now]+way[i].v;
            deep[way[i].y]=deep[now]+1;
            dfs(way[i].y,now);
        }
    cnt[sum[now]+o]--;
}

void solve(int now,int fa) {
    vis[now]=1;
    g[o][0]=1;
    int mx=0;

    for (int i=st[now];i;i=way[i].nxt) 
        if (!vis[way[i].y]) {
            mxdeep=1;
            deep[way[i].y]=1;
            sum[way[i].y]=way[i].v;
            dfs(way[i].y,now);

            mx=max(mx,mxdeep);            //子树中的最大深度 
            ans+=(g[o][0]-1)*f[o][0];
            for (int j=-mxdeep;j<=mxdeep;j++) 
                ans=ans+f[j+o][1]*g[-j+o][0]+f[j+o][0]*g[-j+o][1]+f[j+o][1]*g[-j+o][1];

            for (int j=-mxdeep;j<=mxdeep;j++){
                g[j+o][0]+=f[j+o][0];
                g[j+o][1]+=f[j+o][1];
                f[j+o][0]=f[j+o][1]=0;     //清零 
            }
        }

    for (int i=-mx;i<=mx;i++)
        g[i+o][0]=g[i+o][1]=0;

    for (int i=st[now];i;i=way[i].nxt)
        if (!vis[way[i].y]) {
            sz=size[way[i].y]; root=0;
            findroot(way[i].y,0);
            solve(root,0);
        } 
}

int main() 
{
    scanf("%d",&n);
    for (int i=1;i<n;i++) {
        int u,w,z;
        scanf("%d%d%d",&u,&w,&z);
        add(u,w,(z==1)? 1:-1);
    }
    sz=n; o=n;
    F[0]=N; root=0;
    findroot(1,0);
    solve(root,0);
    printf("%lld",ans);
    return 0;
}