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

Solution

我可能学了假的点分治……
用g[i][0…1]，f[i][0…1]分别表示
前面几个子树以及当前子树路径长度和为i的路径数目
0和1用于区分路径上是否存在前缀和为i的节点(也就是可以设立中转站的节点)
 那么当前子树的贡献就是f[0][0] * g[0][0] + Σf [i][0] * g [-i][1] + f[i][1] * g[-i][0] + f[i][1] * g[-i][1]，
其中i的范围[-d,d]，d为当前子树的深度。

Code

 #include<iostream>
 #include<cstring>
 #include<cstdio>
 #define N (200000+100)
 using namespace std;
 struct node
 {
     int to,next,len;
 } edge[N*];
 int n,k,sum,root,INF;
 long long ans,g[N][],f[N][];
 int t[N*];
 int head[N],num_edge,max_depth;
 int depth[N],d[N],size[N],maxn[N],dis[N];
 bool vis[N];

 void add(int u,int v,int l)
 {
     edge[++num_edge].to=v;
     edge[num_edge].len=l;
     edge[num_edge].next=head[u];
     head[u]=num_edge;
 }

 void Get_root(int x,int fa)
 {
     size[x]=;
     maxn[x]=;
     for (int i=head[x]; i!=; i=edge[i].next)
         if (!vis[edge[i].to] && edge[i].to!=fa)
         {
             Get_root(edge[i].to,x);
             size[x]+=size[edge[i].to];
             maxn[x]=max(maxn[x],size[edge[i].to]);
         }
     maxn[x]=max(maxn[x],sum-size[x]);
     if (maxn[x]<maxn[root]) root=x;
 }

 void Calc(int x,int fa)
 {
     max_depth=max(max_depth,depth[x]);
     if (t[dis[x]]) f[dis[x]][]++;
     else f[dis[x]][]++;
     t[dis[x]]++;
     for (int i=head[x]; i!=; i=edge[i].next)
         if (edge[i].to!=fa && !vis[edge[i].to])
         {
             depth[edge[i].to]=depth[x]+;
             dis[edge[i].to]=dis[x]+edge[i].len;
             Calc(edge[i].to,x);
         }
     t[dis[x]]--;
 }

 void Solve(int x)
 {
     vis[x]=true;
     g[n][]=;//关于这里为什么初始化为1的问题，想了好久，最后还是学姐给我的解答
     //http://blog.csdn.net/wu_tongtong/article/details/79428928
     //可能路径是从当前树根到某一个节点的时候，路径已经平衡，不需要去另一个子树里面找另一个链拼接了
     //所以这一部分的答案也要统计进去。
     int up=;
     for (int i=head[x]; i!=; i=edge[i].next)
         if (!vis[edge[i].to])
         {
             depth[edge[i].to]=;
             dis[edge[i].to]=edge[i].len+n;
             max_depth=;
             Calc(edge[i].to,);
             up=max(up,max_depth);

             ans+=(g[n][]-)*f[n][];
             for (int j=-max_depth; j<=max_depth; ++j)
                 ans+=g[n-j][]*f[n+j][]+g[n-j][]*f[n+j][]+g[n-j][]*f[n+j][];
             for (int j=-max_depth; j<=max_depth; ++j)
             {
                 g[n-j][]+=f[n-j][];
                 g[n-j][]+=f[n-j][];
                 f[n-j][]=f[n-j][]=;
             }
         }
     for (int i=-up; i<=up; ++i)
         g[n-i][]=g[n-i][]=;
     for (int i=head[x]; i!=; i=edge[i].next)
         if (!vis[edge[i].to])
         {
             sum=size[edge[i].to];
             root=;
             Get_root(edge[i].to,);
             Solve(root);
         }
 }

 int main()
 {
     int u,v,l;
     scanf("%d",&n);
     sum=maxn[]=n;
     for (int i=; i<=n-; ++i)
     {
         scanf("%d%d%d",&u,&v,&l);
         add(u,v,l==?-:);
         add(v,u,l==?-:);
     }
     Get_root(,);
     Solve(root);
     printf("%lld",ans);
 }