模板点分治

Posted 2020-12-23 614685877--aakennes

#include<cstdio>
#include<cstring>
#include<iostream>
#include<cmath>
#include<algorithm>
#include<queue>
using namespace std;
const int maxn=1e4+5,INF=0x3f3f3f3f;
int n,k,ans,tot,rt,tsize,cnt,size[maxn],wt[maxn],a[maxn],head[maxn];
bool vis[maxn];
struct Edge{
	int next,to,w;
}e[maxn*2];
inline int read(){
	int s=0,w=1;
	char ch=getchar();
	while(ch<‘0‘||ch>‘9‘){
		if(ch==‘-‘)w=-1;ch=getchar();
	}
	while(ch>=‘0‘&&ch<=‘9‘)s=s*10+ch-‘0‘,ch=getchar();
	return s*w;
	
}
void Add(int x,int y,int z){
	e[++tot].next=head[x];
	e[tot].to=y;
	e[tot].w=z;
	head[x]=tot;
}
void getroot(int u,int fa){
	size[u]=1;wt[u]=0;
	for(int x=head[u];x;x=e[x].next){
		int v=e[x].to;
		if(v==fa||vis[v])continue;
		getroot(v,u);
		size[u]+=size[v];
		wt[u]=max(wt[u],size[v]);
	}
	wt[u]=max(wt[u],tsize-size[u]);
	if(wt[rt]>wt[u])rt=u;
}
void dfs(int u,int fa,int d){
	a[++cnt]=d;
	for(int x=head[u];x;x=e[x].next){
		int v=e[x].to;
		if(v==fa||vis[v])continue;
		dfs(v,u,d+e[x].w);
	}
}
int calc(int u,int d){
	cnt=0;
	dfs(u,0,d);
	sort(a+1,a+cnt+1);
	int sum=0;
	for(int i=1,j=cnt;;i++){
		while(j&&a[i]+a[j]>k)j--;
		if(i>j)break;
		sum+=j-i+1;
	}
	return sum;
}
void DFS(int u){
	ans+=calc(u,0);
	vis[u]=1;
	for(int x=head[u];x;x=e[x].next){
		int v=e[x].to;
		if(vis[v])continue;
		ans-=calc(v,e[x].w);
		rt=0;tsize=size[v];
		getroot(v,0);
		DFS(rt);
	}
}
int main(){
	while(scanf("%d%d",&n,&k)==2){
		ans=0;
		memset(vis,0,sizeof(vis));
		memset(head,0,sizeof(head));
		for(int i=1;i<=n;i++){
			int x=read(),y=read(),z=read();
			Add(x,y,z);Add(y,x,z);
		}
		wt[0]=INF;
		rt=0;tsize=n;getroot(1,0);
		DFS(rt);
		cout<<ans-n<<endl;
	}	
	
	
}