rt,如下程序在跑 n=5×104 时会 TLE,如何优化。
#include<bits/stdc++.h>
#define int long long
using namespace std;
const int N=5e4+5;
bool flag[N];
int n,m,k,u,v,siz[N],cnt[N],tot,dis[N],minn,root;
list<int>edge[N];
void build(int u,int v)
{
edge[u].push_back(v);
edge[v].push_back(u);
}
void getroot(int u,int fa,int treesiz)
{
siz[u]=1;
int maxn=treesiz-siz[u];
for(auto v:edge[u])
{
if(v==fa||flag[v])
continue;
getroot(v,u,treesiz);
siz[u]+=siz[v];
maxn=max(maxn,siz[v]);
}
if(maxn<minn)root=u,minn=maxn;
}
void getdis(int u,int fa,int dist)
{
siz[u]=1;
dis[++tot]=dist;
for(auto v:edge[u])
{
if(v==fa||flag[v])
continue;
getdis(v,u,dist+1);
siz[u]+=siz[v];
}
}
int DFS(int u,int fa,int treesiz)
{
int res=0;
root=u,minn=treesiz-1;
getroot(u,0,treesiz);
u=root,flag[u]=true;
memset(cnt,0,sizeof cnt);
cnt[0]=1;
for(auto v:edge[u])
{
if(v==fa||flag[v])
continue;
tot=0;
getdis(v,u,1);
for(int i=1;i<=tot;i++)
if(k-dis[i]>=0)res+=cnt[k-dis[i]];
for(int i=1;i<=tot;i++)
cnt[dis[i]]++;
}
for(auto v:edge[u])
{
if(v==fa||flag[v])
continue;
res+=DFS(v,u,siz[v]);
}
return res;
}
signed main()
{
cin>>n>>k;
for(int i=1;i<n;i++)
cin>>u>>v,build(u,v);
cout<<DFS(1,0,n);
return 0;
}