理论参考https://wenku.baidu.com/view/8861df38376baf1ffc4fada8.html?re=view
树的分治
#include<cstdio>
#include<cstring>
#include<algorithm>
using namespace std;
const int maxn=10010;
int N,K;
int root; //根结点
int ans; //结果
int Max;
struct node
{
int v;
int next;
int w;
}edge[maxn*2]; //无向图
int tot; //结点总数total
int head[maxn]; //头结点数组head
int size[maxn]; //以v为根的子树节点数size
int maxv[maxn];
int vis[maxn]; //遍历标志visit
int dis[maxn]; //距离数组distance
int num;
void init()
{
tot=0;
ans=0;
memset(head,-1,sizeof(head));
memset(vis,0,sizeof(vis));
}
//建立邻接表
void add_edge(int u,int v,int w)
{
edge[tot].v=v;
edge[tot].w=w;
edge[tot].next=head[u];//形成了链表,以-1结尾,以head为头结点,每个结点都是结点u的孩子。
head[u]=tot++;
}
//计算子树所含节点数
void dfssize(int u,int f) //除去f点外的规模,f一般为父结点
{
size[u]=1; //本身 一个结点
maxv[u]=0; //其最大子树的结点数
for(int i=head[u];i!=-1;i=edge[i].next)
{
int v=edge[i].v;
if(v==f||vis[v])
continue;
dfssize(v,u);
size[u]+=size[v];
if(size[v]>maxv[u])
maxv[u]=size[v];
}
}
//确定树的重心,即最佳的根结点。
void dfsroot(int r,int u,int f) //r为根结点,u为比较结点
{
if(size[r]-size[u]>maxv[u]) //size[r]-size[u]是u上面部分的树的尺寸,maxv[u]是u的最大子树的结点数。
maxv[u]=size[r]-size[u];//此处将树看作二叉树,把最大子树作为一枝,剩余的看作一枝 。将最大的一枝存入maxv中
if(maxv[u]<Max) Max=maxv[u],root=u; //MAX是maxv中最小值,由于maxv必不小于总结点的一半 ,故MAX存储的是当前最平衡结点的最大枝个数
for(int i=head[u];i!=-1;i=edge[i].next)
{
int v=edge[i].v;
if(v==f||vis[v])
continue;
dfsroot(r,v,u);
}
}
//求每个点离重心的距离
void dfsdis(int u,int d,int f)
{
dis[num++]=d;
for(int i=head[u];i!=-1;i=edge[i].next)
{
int v=edge[i].v;
if(v!=f&&!vis[v])
dfsdis(v,d+edge[i].w,u);
}
}
//计算以u为根的子树中有多少点对的距离小于等于K
int calc(int u,int d) //d为通向u的路径长度
{
int ret=0; //点对的个数
num=0; //初始化
dfsdis(u,d,0);
sort(dis,dis+num);
int i=0,j=num-1; //0点为根结点本身
while(i<j)
{
while(dis[i]+dis[j]>K&&i<j)j--;
ret+=j-i;
i++;
}
return ret;
}
void dfs(int u)
{
Max=N;
dfssize(u,0);
dfsroot(u,u,0); //设置u为根
ans+=calc(root,0); //u树中的点对
vis[root]=1; //根结点已遍历
for(int i=head[root];i!=-1;i=edge[i].next)
{
int v=edge[i].v;
if(!vis[v])
{
ans-=calc(v,edge[i].w); //删除掉同属于v结点且满足 dis[i]+dis[j]<K的点对。
dfs(v); //加上以v为结点仅可过v结点满足 dis[i]+dis[j]<K的点对。
}
}
}
int main()
{
while(scanf("%d%d",&N,&K)!=EOF)
{
if(!N&&!K)
break;
int u,v,w;
init();
for(int i=1;i<N;i++)
{
scanf("%d%d%d",&u,&v,&w);
add_edge(u,v,w);
add_edge(v,u,w);
}
dfs(1);
printf("%d\n",ans);
}
return 0;
}
|