初学点分治TLE on 26 求调

#include<bits/stdc++.h>
using namespace std;
const int N=5e4+100;
int n,k,idx,root,siz[N],bcs,head[N];
int L,p,dis[N],cnt[N];
struct edge{
	int v,next;
}e[N*2];
void con(int u,int v){
	e[++idx].v=v;
	e[idx].next=head[u];
	head[u]=idx;
}
bool vis[N];
long long ans;
void findroot(int u,int fa,int nod){
	siz[u]=1;
	int ms=0;
	for(int i=head[u];i;i=e[i].next){
		int v=e[i].v;
		if(v==fa||vis[v]) continue;
		findroot(v,u,nod);
		siz[u]+=siz[v];
		ms=max(ms,siz[u]);
	}
	ms=max(ms,nod-siz[u]);
	if(ms<bcs) {
		root=u;
		bcs=ms;
	}
}
void get_dis(int u,int fa,int d){
	if(d>k) return;
	dis[++p]=d;
	siz[u]=1;
	for(int i=head[u];i;i=e[i].next){
		int v=e[i].v;
		if(v==fa||vis[v]) continue;
		get_dis(v,u,d+1);
		siz[u]+=siz[v];
	}
}
void lc(int u,int fa,int nod){
	root=u;
	bcs=nod-1;
	findroot(u,fa,nod);
	p=L=1;
	dis[1]=0;
	cnt[0]=1;
	vis[root]=1;
	for(int i=head[root];i;i=e[i].next){
		int v=e[i].v;
		if(v==fa||vis[v]) continue;
		get_dis(v,root,1);
		for(int j=L+1;j<=p;j++) {
			if(k-dis[j]>=0) {
				ans+=cnt[k-dis[j]];
			}
		}
		for(int j=L+1;j<=p;j++) {
			cnt[dis[j]]++;
		}
		L=p;
	}
	for(int i=1;i<=p;i++) cnt[dis[i]]--;
	for(int i=head[root];i;i=e[i].next){
		int v=e[i].v;
		if(v==fa||vis[v]) continue;
		lc(v,root,siz[v]);
	}
}
signed main(){
	scanf("%d%d",&n,&k);
	for(int i=1;i<n;i++) {
		int u,v;
		scanf("%d%d",&u,&v);
		con(u,v);
		con(v,u);
	}
	lc(1,0,n);
	printf("%lld",ans);
	
	return 0;
}