Bzoj3697:采药人的路径:树的点分治

题目链接：3697:采药人的路径

Orz黄学长

#include<cstdio>
#include<cstdlib>
#include<iostream>
#include<algorithm>
#define LL long long
using namespace std;
const int maxn=200010;
int n,tot=1,h[maxn],s[maxn],dp[maxn];
struct edge{int to,next,w;}G[maxn*2];
LL t[maxn],g[maxn][2],f[maxn][2],ans=0;
int size,root,mxd,dep[maxn],dis[maxn];
bool flag[maxn];

void add(int x,int y,int z){
	G[++tot].to=y;G[tot].next=h[x];h[x]=tot;G[tot].w=z;
	G[++tot].to=x;G[tot].next=h[y];h[y]=tot;G[tot].w=z;
}

void getroot(int x,int fa){
	s[x]=1; dp[x]=0;
	for (int i=h[x];i;i=G[i].next){
		int v=G[i].to;
		if (v==fa||flag[v]) continue;
		getroot(v,x); s[x]+=s[v];
		dp[x]=max(dp[x],s[v]);
	}dp[x]=max(dp[x],size-s[x]);
	if (dp[x]<dp[root]) root=x;
}

void getdep(int x,int fa){
	mxd=max(mxd,dep[x]);
	if (t[dis[x]]) f[dis[x]][1]++;
	else f[dis[x]][0]++; t[dis[x]]++;
	for (int i=h[x];i;i=G[i].next){
		int v=G[i].to;
		if (v==fa||flag[v]) continue;
		dep[v]=dep[x]+1;
		dis[v]=dis[x]+G[i].w;
		getdep(v,x);
	}t[dis[x]]--;
}

void work(int x){
	int mx=0; flag[x]=1; g[n][0]=1;
	for (int i=h[x];i;i=G[i].next){
		int v=G[i].to;
		if (flag[v]) continue;
		dep[v]=1; dis[v]=n+G[i].w;
		mxd=1; getdep(v,0);
		mx=max(mx,mxd);
		ans+=(g[n][0]-1)*f[n][0];
		for (int j=-mxd;j<=mxd;++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-mxd;j<=n+mxd;++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;
	for (int i=h[x];i;i=G[i].next){
		int v=G[i].to;
		if (flag[v]) continue;
		dp[0]=size=s[v];
		getroot(v,root=0);
		work(root);
	}
}

int main(){
	scanf("%d",&n);
	for (int i=1;i<n;++i){
		int x,y,z;
		scanf("%d%d%d",&x,&y,&z);
		if (z==0) z=-1;	add(x,y,z);
 	}
	dp[0]=size=n; getroot(1,root=0);
	ans=0; work(root); printf("%lld",ans);
}