CF161D Distance in Tree |点分治

题目大意

输入点数为N一棵树

求树上长度恰好为K的路径个数

输入格式

第一行两个数字N,K,如题意

接下来的N−1行中,每行两个整数u,v表示一条树边(u,v)

输出格式

一个整数ans,如题意

---

直接套模板，没啥好说的

#include<cstdio>
#include<algorithm>
#include<cstring>
#include<cmath>
#include<iostream>
#include<map>
using namespace std;
inline int read(){
	int x=0,f=1; char c=getchar();
	while(c<'0'||c>'9'){ if(c=='-')f=-1; c=getchar(); }
	while('0'<=c&&c<='9'){ x=x*10+c-'0'; c=getchar(); }
	return x*f;
}
const int N=5e4+10;
int nxt[N<<1],head[N],go[N<<1],tot;
inline void add(int u,int v){
	nxt[++tot]=head[u]; head[u]=tot; go[tot]=v;
	nxt[++tot]=head[v]; head[v]=tot; go[tot]=u;	
}
bool vis[N];
int n,k,size[N],sum,rt,maxp;
inline void getrt(int u,int fa){
	size[u]=1; int Max=0;
	for(int i=head[u];i;i=nxt[i]){
		int v=go[i];
		if(v==fa||vis[v])continue;
		size[u]+=size[v];
		Max=max(Max,size[v]);
	}
	Max=max(Max,sum-size[u]);
	if(Max<maxp)maxp=Max,rt=u;
}
int judge[N],rem[N],dis[N];
inline void getdis(int u,int fa){
	if(dis[u]>k)return;
	rem[++rem[0]]=dis[u];
	for(int i=head[u];i;i=nxt[i]){
		int v=go[i];
		if(v==fa||vis[v])continue;
		dis[v]=dis[u]+1;
		getdis(v,u);
	}
}
long long ans=0;
inline void calc(int u){
	for(int i=head[u];i;i=nxt[i]){
		int v=go[i];
		if(vis[v])continue;
		rem[0]=0; dis[v]=1; getdis(v,u);
		for(int i=1;i<=rem[0];i++)ans+=judge[k-rem[i]];
		for(int i=1;i<=rem[0];i++)judge[rem[i]]++;
	}
	for(int i=0;i<=n;i++){
		if(!judge[i])break;
		judge[i]=0;
	}
}
inline void solve(int u){
	judge[0]=vis[u]=1; calc(u);
	for(int i=head[u];i;i=nxt[i]){
		int v=go[i];
		if(vis[v])continue;
		sum=maxp=size[v]; getrt(v,0);
		solve(rt);
	}
}
signed main(){
	n=read(),k=read();
	for(int i=1;i<n;i++)add(read(),read());
	sum=maxp=n; getrt(1,0);
	solve(rt);
	cout<<(ans>>1)<<endl;
}