BZOJ.1468.Tree(点分治)

BZOJ1468
POJ1741

题意: 计算树上距离<=K的点对数

我们知道树上一条路径要么经过根节点，要么在同一棵子树中。
于是对一个点x我们可以这样统计: 计算出所有点到它的距离dep[]，排序后可以O(n)求得<=K的点对数量。
但画个图后我们可以发现，对于在同一棵子树中的路径被重复计算过了。于是我们Ans-=Calc(v)，减去一棵子树中的路径答案，但是这并不是之前x到它们的路径，于是给v的dep[]设一个初始值为w(x->v路径权值)。
这样x的答案就计算完了，将这一过程记作Solve(x)。
考虑如何计算所有点。DFS的效率是和树深有关的。计算x->v时，我们选取v子树上的重心作为下次Solve()的参数。
选取重心每次都会使树的节点个数减半，因此递归深度最坏(链)是logn的。
这样便可O(nlog^2n)解决该题。

//2524kb	656ms
#include <cstdio>
#include <cctype>
#include <algorithm>
//#define gc() getchar()
#define gc() (SS==TT&&(TT=(SS=IN)+fread(IN,1,MAXIN,stdin),SS==TT)?EOF:*SS++)
const int N=40005,MAXIN=1e5;

int n,K,root,Min,Ans,Enum,H[N],nxt[N<<1],to[N<<1],val[N<<1],dep[N],sz[N],D[N];
bool vis[N];
char IN[MAXIN],*SS=IN,*TT=IN;

inline int read()
{
	int now=0;register char c=gc();
	for(;!isdigit(c);c=gc());
	for(;isdigit(c);now=now*10+c-'0',c=gc());
	return now;
}
inline void AddEdge(int u,int v,int w)
{
	to[++Enum]=v, nxt[Enum]=H[u], H[u]=Enum, val[Enum]=w;
	to[++Enum]=u, nxt[Enum]=H[v], H[v]=Enum, val[Enum]=w;	
}
void Get_Root(int x,int f,int tot)
{
	int mx=0; sz[x]=1;
	for(int v,i=H[x]; i; i=nxt[i])
		if(!vis[to[i]] && to[i]!=f)
		{
			Get_Root(v=to[i],x,tot), sz[x]+=sz[v];
			if(sz[v]>mx) mx=sz[v];
		}
	mx=std::max(mx,tot-sz[x]);
	if(mx<Min) Min=mx, root=x;
}
void DFS(int x,int f)
{
	D[++D[0]]=dep[x];
	for(int i=H[x]; i; i=nxt[i])
		if(!vis[to[i]]&&to[i]!=f)
			dep[to[i]]=dep[x]+val[i], DFS(to[i],x);
}
int Calc(int x,int v)
{
	D[0]=0, dep[x]=v, DFS(x,-1);
	std::sort(D+1,D+1+D[0]);
	int l=1,r=D[0],res=0;
	while(l<r)
		if(D[l]+D[r]<=K) res+=r-l,++l;
		else --r;
	return res;
}
void Solve(int x)
{
	Ans+=Calc(x,0), vis[x]=1;
	for(int v,i=H[x]; i; i=nxt[i])
		if(!vis[to[i]])
		{
			Ans-=Calc(v=to[i],val[i]);
			Min=N, Get_Root(v,x,sz[v]), Solve(root);
		}
}

int main()
{
	n=read();
	for(int u,v,w,i=1; i<n; ++i) u=read(),v=read(),w=read(),AddEdge(u,v,w);
	K=read();
	Min=N, Get_Root(1,-1,n), Solve(root);
	printf("%d",Ans);

	return 0;
}