luogu P4178 Tree

背景：

$hehe...$

题目传送门：

https://www.luogu.org/problemnew/show/P4178

题意：

一棵树，问树上最短距离不大于$k$的点对数。

思路：

点分治模板吧。
只是统计答案时将$dis$升序，然后维护一个单调队列统计答案（$di{s}_i$递增，那么另一个递减）。

代码：

#include<cstdio>
#include<cstring>
#include<algorithm>
using namespace std;
	int n,m,len=0,ans=0;
	int last[200010],size[200010],msize[200010],dis[200010];
	bool bz[200010];
	struct node{int x,y,z,next;} a[200010];
int SIZE,MIN,ROOT,tot;
void ins(int x,int y,int z)
{
	a[++len]=(node){x,y,z,last[x]}; last[x]=len;
}
void find_root(int x,int fa)
{
	size[x]=1;
	msize[x]=0;
	for(int i=last[x];i;i=a[i].next)
	{
		int y=a[i].y;
		if(y==fa||bz[y]) continue;
		find_root(y,x);
		size[x]+=size[y];
		msize[x]=max(msize[x],size[y]);
	}
	msize[x]=max(msize[x],SIZE-size[x]);
	if(MIN>msize[x]) MIN=msize[x],ROOT=x;
}
void get_dis(int x,int fa,int z)
{
	dis[++tot]=z;
	for(int i=last[x];i;i=a[i].next)
	{
		int y=a[i].y;
		if(bz[y]||y==fa) continue;
		get_dis(y,x,z+a[i].z);
	}
}
void solve(int x,int z,int op)
{
	tot=0;
	get_dis(x,0,z);
	sort(dis+1,dis+tot+1);
	int tail=tot;
	for(int i=1;i<=tot;i++)
	{
		if(dis[i]>m) break;
		while(i<tail&&dis[i]+dis[tail]>m) tail--;
		if(i==tail) break;
		ans+=op*(tail-i+1-1);
	}
}
void dfs(int x)
{
	bz[x]=true;
	solve(x,0,1);
	for(int i=last[x];i;i=a[i].next)
	{
		int y=a[i].y;
		if(bz[y]) continue;
		solve(y,a[i].z,-1);
		SIZE=size[y],MIN=SIZE,ROOT=0;
		find_root(y,0);
		dfs(ROOT);
	}
}
int main()
{
	int x,y,z;
	scanf("%d",&n);
	for(int i=1;i<n;i++)
	{
		scanf("%d %d %d",&x,&y,&z);
		ins(x,y,z),ins(y,x,z);
	}
	scanf("%d",&m);
	SIZE=n,MIN=n,ROOT=0;
	find_root(1,0);
	dfs(ROOT);
	printf("%d",ans);
}