POJ 1741 Tree | 树分治

求树上距离小于等于K的点对对数

---

#include<cstdio>
#include<algorithm>
#include<cstring>
#define N 10005
typedef long long ll;
using namespace std;
int n,K,fa[N],sze[N],son[N],dis[N],head[N];
int ecnt;
bool vis[N];
ll ans;
struct edge
{
    int nxt,v,w;
}e[2*N];
void add(int u,int v,int w)//加边
{
    e[++ecnt].v=v,e[ecnt].w=w,e[ecnt].nxt=head[u],head[u]=ecnt;
    e[++ecnt].v=u,e[ecnt].w=w,e[ecnt].nxt=head[v],head[v]=ecnt;
}
void init()//初始化
{
    ans=ecnt=0;
    for (int i=1;i<=n;i++)
	head[i]=vis[i]=0;
}
int calcG(int sv)//计算重心
{
    static int qn,que[N];
    int u,v,mx=n,G; 
    que[qn=1]=sv,fa[sv]=0;
    for (int ql=1;ql<=qn;ql++)//以sv为根bfs求父子关系
    {
	sze[u=que[ql]]=1,son[u]=0;
	for (int i=head[u];i;i=e[i].nxt)
	{
	    if (vis[v=e[i].v] || v==fa[u]) continue;
	    fa[v]=u,que[++qn]=v;
	}
    }
    for (int ql=qn;ql>=1;ql--)//bfs的逆序更新每个节点的最大子树大小
    {
    //son数组表示每个节点为根最大子树大小
	u=que[ql],v=fa[u];
	if (qn-sze[u]>son[u]) son[u]=qn-sze[u];//如果他的子树比qn/2要小,就更新成qn-sze  
	if (son[u]<mx) G=u,mx=son[u];//更新重心
	if (!v) break;
	sze[v]+=sze[u];//更新爸爸的子树
	if (sze[u]>son[v]) son[v]=sze[u];
    }
    return G;
}
inline ll calc(int sv,int L)
{
    static int qn,que[N],d[N];
    int u,v,d_n=0;
    que[qn=1]=sv,dis[sv]=L,fa[sv]=0;
    for (int ql=1;ql<=qn;ql++)
    {
	d[d_n++]=dis[u=que[ql]];
	for (int i=head[u];i;i=e[i].nxt)
	{
	    if (vis[v=e[i].v] || v==fa[u]) continue ;
	    fa[v]=u,dis[v]=dis[u]+e[i].w,que[++qn]=v;
	}
    }
    ll cnt=0;
    sort(d,d+d_n);
    int l=0,r=d_n-1;
    while (l<r)
    {
	if (d[l]+d[r]<=K) cnt+=r-l++;
	else --r;
    }
    return cnt;
}
void solve(int u)
{
    int G=calcG(u);//求以u为根的子树的重心
    vis[G]=1;
    ans+=calc(G,0);//加上当前所在树中的答案
	if (!vis[e[i].v]) ans-=calc(e[i].v,e[i].w);//减去每个子树的答案,因为子树的答案在子树中已经计算过了,不能重复计算
    for (int i=head[G];i;i=e[i].nxt)
	if (!vis[e[i].v]) solve(e[i].v);//递归每个子树
}
int main()
{
    while (scanf("%d%d",&n,&K),n || K)
    {
	init();
	for (int i=1;i<n;i++)
	{
	    int u,v,w;
	    scanf("%d%d%d",&u,&v,&w);
	    add(u,v,w);
	}
	solve(1);//以1为根
	printf("%lld\n",ans);
    }
    return 0;
}