题意:求树上距离小于等于m的点对有多少个。
#include<iostream>
#include<cstdio>
#include<cstring>
#include<algorithm>
#include<vector>
using namespace std;
const int maxn=1e4+100;
struct node{
int x,w;
node(int _x,int _w)
{
x=_x;
w=_w;
}
};
vector<node>a[maxn];
int hvy,n,m,vis[maxn],siz[maxn],ma[maxn];
int S; //树的大小
void dfs_hvy(int u,int fa)//找重心
{
siz[u]=1;
ma[u]=0;
for(int i=0;i<a[u].size();i++)
{
int v=a[u][i].x;
if(v==fa||vis[v]) continue;
dfs_hvy(v,u);
siz[u]+=siz[v];
ma[u]=max(ma[u],siz[v]);
}
ma[u]=max(ma[u],S-siz[u]);
if(!hvy||ma[u]<ma[hvy])// 找重心,以重心为根的树的最大子树最小
hvy=u;
}
int dep[maxn],tmp[maxn];
void dfs_dep(int u,int fa)
{
tmp[++tmp[0]]=dep[u];
for(int i=0;i<a[u].size();i++)
{
int v=a[u][i].x,w=a[u][i].w;
if(v==fa||vis[v]) continue;///注意找的是未访问的
dep[v]=dep[u]+w;
dfs_dep(v,u);
}
}
int dfs_sum(int u,int d)
{
tmp[0]=0;
int res=0;
dep[u]=d;
dfs_dep(u,0);//找出孩子节点到u的距离
sort(tmp+1,tmp+tmp[0]+1);///
int j=1;
for(int i=1,j=tmp[0];i<=tmp[0];i++)
{
while(j>i&&tmp[i]+tmp[j]>m)j--;///
res+=max(j-i,0);///
}
return res;
}
int ans;
void dfs_ans(int u)
{
vis[u]=1;
ans+=dfs_sum(u,0);//加上经过重心的点对数(加的里面包括了不经过的,需要在下面再减去)
for(int i=0;i<a[u].size();i++)
{
int v=a[u][i].x,w=a[u][i].w;
if(vis[v]) continue;
ans-=dfs_sum(v,w);/不经过的在直接孩子处减去
S=siz[v],hvy=0;
dfs_hvy(v,u);
dfs_ans(hvy);
}
}
int main()
{
while(scanf("%d%d",&n,&m),n+m)
{
for(int i=1;i<=n;i++)
a[i].clear();
memset(vis,0,sizeof vis);
for(int i=1;i<n;i++)
{
int x,y,w;
scanf("%d%d%d",&x,&y,&w);
a[x].push_back(node(y,w));
a[y].push_back(node(x,w));
}
S=n;hvy=0;
dfs_hvy(1,0);//找重心
ans=0;
dfs_ans(hvy);
printf("%d\n",ans);
}
return 0;;;
}
从这篇博客学到的