题意:输入一棵树,输出前k小的点对最短距离dis(i,j)的和。
模拟,官方题解说得很清楚了。不重复了。
需要注意的是,复杂度要O(n+k),不能用set,map之类的标记是否访问。
一开始TLE了,去掉标记后wa了。最后发现对队列的元素加个前缀,就可以了,即标记该条边是从哪个点延伸的。
#include <cstdio>
#include <cstring>
#include <iostream>
#include <cmath>
#include <algorithm>
#include <vector>
#include <string>
#include <set>
#include <queue>
#include <map>
using namespace std;
#define ll long long
#define MP make_pair
#define inf 1e9
#define eps 1e-8
#define maxn 110000
#define maxm 210000
int que[2100000][4];
struct Edge{
int v,nxt;
}e[maxm];
int head[maxn],esz;
void init(){esz=0;memset(head,-1,sizeof(head));}
void addedge(int u,int v){
e[esz].v=v,e[esz].nxt=head[u];
head[u]=esz++;
}
int main(){
int t;
scanf("%d",&t);
while(t--){
int n,m;
scanf("%d%d",&n,&m);
m<<=1;
int sz=0;
init();
for(int i=1;i<n;++i){
int u,v;
scanf("%d%d",&u,&v);
addedge(u,v);
addedge(v,u);
if(sz<m){
que[sz][0]=u,que[sz][1]=v,que[sz][2]=1;
que[sz][3]=-1;
++sz;
}
if(sz<m){
que[sz][0]=v,que[sz][1]=u,que[sz][2]=1;
que[sz][3]=-1;
++sz;
}
}
int H=0,T=sz;
ll ans=0;
for(int i=0;i<m;++i){
int dis = que[H][2];
int u = que[H][0], v = que[H][1];
int pre = que[H][3];
++H;
for(int j=head[u];j!=-1 && T<m;j=e[j].nxt){
int uu = e[j].v;
if(uu==v || uu==pre) continue;
que[T][0]=uu,que[T][1]=v,que[T][2]=dis+1;
que[T][3]=u;
T++;
}
ans+=dis;
}
printf("%I64d\n",ans>>1);
}
return 0;
}