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

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

  1 #include<iostream>
  2 #include<cstring>
  3 #include<cstdio>
  4 #define N (200000+100)
  5 using namespace std;
  6 struct node
  7 {
  8     int to,next,len;
  9 } edge[N*2];
 10 int n,k,sum,root,INF;
 11 long long ans,g[N][2],f[N][2];
 12 int t[N*2];
 13 int head[N],num_edge,max_depth;
 14 int depth[N],d[N],size[N],maxn[N],dis[N];
 15 bool vis[N];
 16 
 17 void add(int u,int v,int l)
 18 {
 19     edge[++num_edge].to=v;
 20     edge[num_edge].len=l;
 21     edge[num_edge].next=head[u];
 22     head[u]=num_edge;
 23 }
 24 
 25 void Get_root(int x,int fa)
 26 {
 27     size[x]=1;
 28     maxn[x]=0;
 29     for (int i=head[x]; i!=0; i=edge[i].next)
 30         if (!vis[edge[i].to] && edge[i].to!=fa)
 31         {
 32             Get_root(edge[i].to,x);
 33             size[x]+=size[edge[i].to];
 34             maxn[x]=max(maxn[x],size[edge[i].to]);
 35         }
 36     maxn[x]=max(maxn[x],sum-size[x]);
 37     if (maxn[x]<maxn[root]) root=x;
 38 }
 39 
 40 void Calc(int x,int fa)
 41 {
 42     max_depth=max(max_depth,depth[x]);
 43     if (t[dis[x]]) f[dis[x]][1]++;
 44     else f[dis[x]][0]++;
 45     t[dis[x]]++;
 46     for (int i=head[x]; i!=0; i=edge[i].next)
 47         if (edge[i].to!=fa && !vis[edge[i].to])
 48         {
 49             depth[edge[i].to]=depth[x]+1;
 50             dis[edge[i].to]=dis[x]+edge[i].len;
 51             Calc(edge[i].to,x);
 52         }
 53     t[dis[x]]--;
 54 }
 55 
 56 void Solve(int x)
 57 {
 58     vis[x]=true;
 59     g[n][0]=1;//关于这里为什么初始化为1的问题，想了好久，最后还是学姐给我的解答
 60     //http://blog.csdn.net/wu_tongtong/article/details/79428928
 61     //可能路径是从当前树根到某一个节点的时候，路径已经平衡，不需要去另一个子树里面找另一个链拼接了
 62     //所以这一部分的答案也要统计进去。
 63     int up=0;
 64     for (int i=head[x]; i!=0; i=edge[i].next)
 65         if (!vis[edge[i].to])
 66         {
 67             depth[edge[i].to]=1;
 68             dis[edge[i].to]=edge[i].len+n;
 69             max_depth=1;
 70             Calc(edge[i].to,0);
 71             up=max(up,max_depth);
 72 
 73             ans+=(g[n][0]-1)*f[n][0];
 74             for (int j=-max_depth; j<=max_depth; ++j)
 75                 ans+=g[n-j][0]*f[n+j][1]+g[n-j][1]*f[n+j][0]+g[n-j][1]*f[n+j][1];
 76             for (int j=-max_depth; j<=max_depth; ++j)
 77             {
 78                 g[n-j][0]+=f[n-j][0];
 79                 g[n-j][1]+=f[n-j][1];
 80                 f[n-j][0]=f[n-j][1]=0;
 81             }
 82         }
 83     for (int i=-up; i<=up; ++i)
 84         g[n-i][0]=g[n-i][1]=0;
 85     for (int i=head[x]; i!=0; i=edge[i].next)
 86         if (!vis[edge[i].to])
 87         {
 88             sum=size[edge[i].to];
 89             root=0;
 90             Get_root(edge[i].to,0);
 91             Solve(root);
 92         }
 93 }
 94 
 95 int main()
 96 {
 97     int u,v,l;
 98     scanf("%d",&n);
 99     sum=maxn[0]=n;
100     for (int i=1; i<=n-1; ++i)
101     {
102         scanf("%d%d%d",&u,&v,&l);
103         add(u,v,l==0?-1:1);
104         add(v,u,l==0?-1:1);
105     }
106     Get_root(1,0);
107     Solve(root);
108     printf("%lld",ans);
109 }