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

题目大意：给出一棵n个结点的树，边权为1或-1，求有多少条路径边权和为0并且能在路径上找到一个在起点与终点之间的休息点(不能是起点或终点)满足起点到休息点、休息点到终点的路径边权为0

不是那么裸的点分治。

首先考虑在一棵子树中如何找到边权和为0的路径。由于边权只有1和-1，和的绝对值不会超过n，记录一下权值为i的路径分别有多少条，其中i属于[-n,n]。这样i与-i组合形成的路径边权和为0。

再考虑找休息点。在子树中由根拼接两条路径组成的一整条路径中，休息点可能在三个位置：起点所在的半条路径，终点所在的半条路径，根。

这样的话记录两个数组：f(i,j)表示之前扫过的子树中权值为i的路径的条数,j代表在根与终点之间是否有某个点满足从根到这个点的权值为i；g(i,j)表示当前的子树，其余与f的描述相同。

则统计答案时枚举i，分别按照休息点在起点所在的半条路径( g(i,1) * f(-i,0) )，休息点在终点所在的半条路径( g(i,0) * f(-i,1) ) ，休息点在两端均可( g(i,1) * f(-i,1) )，休息点在根( (g(0,0)-1) * f(0,0) )统计答案即可。

#include <cstdio>
#include <cstring>
#include <algorithm>
#define N 100005
using namespace std;
typedef long long LL;
struct Edge {
int from,to,nxt,val;
Edge() {}
Edge(int _from,int _to,int _nxt,int _val):
from(_from),to(_to),nxt(_nxt),val(_val) {}
}e[N*2];
int n,tot=-1,fir
;
void Add_Edge(int x,int y,int v) {
e[++tot]=Edge(x,y,fir[x],v), fir[x]=tot;
e[++tot]=Edge(y,x,fir[y],v), fir[y]=tot;
return ;
}
int sum,root,max_dpt,dpt
,dist
,pa
,max_siz
,siz
,t[N*2];
bool vis
;
LL ans,f[N*2][2],g[N*2][2];
void dfs_root(int x,int from) {
pa[x]=from, siz[x]=1, max_siz[x]=0;
for(int i=fir[x];~i;i=e[i].nxt) {
if(e[i].to==from || vis[e[i].to]) continue;
dfs_root(e[i].to,x);
siz[x]+=siz[e[i].to];
max_siz[x]=max(max_siz[x],siz[e[i].to]);
}
max_siz[x]=max(max_siz[x],sum-siz[x]);
if(max_siz[x]<max_siz[root]) root=x;
return ;
}
void get_root(int x) {
root=0;
sum=siz[x];
dfs_root(x,0);
siz[pa[root]]=sum-siz[root];
return ;
}
void dfs_dpt(int x,int from) {
max_dpt=max(max_dpt,dpt[x]);
if(t[dist[x]]) f[dist[x]][1]++;
else f[dist[x]][0]++;
t[dist[x]]++;
for(int i=fir[x];~i;i=e[i].nxt) {
if(e[i].to==from || vis[e[i].to]) continue;
dpt[e[i].to]=dpt[x]+1;
dist[e[i].to]=dist[x]+e[i].val;
dfs_dpt(e[i].to,x);
}
t[dist[x]]--;
return ;
}
void calc(int x) {
int maxx=0;
g
[0]=1;
for(int i=fir[x];~i;i=e[i].nxt) {
if(vis[e[i].to]) continue;
dist[e[i].to]=n+e[i].val;
dpt[e[i].to]=1;
max_dpt=1;
dfs_dpt(e[i].to,0);
maxx=max(max_dpt,maxx);
ans+=(g
[0]-1)*f
[0];
for(int j=-max_dpt;j<=max_dpt;j++)
ans+=g[n-j][1]*f[n+j][1]+g[n-j][0]*f[n+j][1]+g[n-j][1]*f[n+j][0];
for(int j=n-max_dpt;j<=n+max_dpt;j++)
g[j][0]+=f[j][0], g[j][1]+=f[j][1], f[j][0]=f[j][1]=0;
}
for(int i=n-maxx;i<=n+maxx;i++) g[i][0]=g[i][1]=0;
return ;
}
void Divide(int x) {
vis[x]=true;
calc(x);
for(int i=fir[x];~i;i=e[i].nxt) {
if(vis[e[i].to]) continue;
get_root(e[i].to);
Divide(root);
}
return ;
}
int main() {
memset(fir,-1,sizeof fir);
scanf("%d",&n);
for(int i=1;i<n;i++) {
int x,y,z;
scanf("%d%d%d",&x,&y,&z);
if(!z) z=-1;
Add_Edge(x,y,z);
}
max_siz[0]=siz[1]=n;
get_root(1);
Divide(root);
printf("%lld\n",ans);
return 0;
}