hdu 5102 樹上前k短路徑長度和

http://acm.hdu.edu.cn/showproblem.php?pid=5102

給一棵樹，求出所有節點的距離中前k小的路徑長度和

由於路徑長度的定義爲兩點之間的邊的個數，所有遍歷1~n-1條邊組成的路徑，暴力擦線過，3000+ms

#include <cstdio>
#include <cstdlib>
#include <cmath>
#include <cstring>
#include <string>
#include <queue>
#include <map>
#include <iostream>
#include <sstream>
#include <algorithm>
using namespace std;
#define RD(x) scanf("%d",&x)
#define RD2(x,y) scanf("%d%d",&x,&y)
#define RD3(x,y,z) scanf("%d%d%d",&x,&y,&z)
#define clr0(x) memset(x,0,sizeof(x))
#define clr1(x) memset(x,-1,sizeof(x))
#define eps 1e-9
const double pi = acos(-1.0);
typedef long long LL;
const int inf = 0x7fffffff;
const int maxn = 1e5+5;
LL sum;
vector<int> g[maxn];
int n,k,goal;
void dfs(int u,int fa,int len)
{
    if(len == goal){
        sum++;
        return ;
    }
    for(int i = 0;i < g[u].size();++i){
        int v = g[u][i];
        if(v != fa)
            dfs(v,u,len+1);
    }
}
int main()
{
    int _;
    RD(_);
    while(_--){
        int u,v;
        RD2(n,k);
        if(k < n){
            while(--n)
                RD2(u,v);
            printf("%d\n",k);
            continue;
        }
        for(int i = 1;i <= n;++i)
            g[i].clear();
        int __ = n;
        while(--__){
            RD2(u,v);
            g[u].push_back(v);
            g[v].push_back(u);
        }
        k -= (n - 1);
        LL ans = n - 1,cnt;
        for(int l = 2;l < n;++l){
            cnt = 0,goal = l;
            for(int j = 1;j <= n;++j){
                sum = 0;
                dfs(j,-1,0);
                cnt += sum;
                if(cnt/2 >= k)
                    break;
            }
            cnt/=2;
            if(cnt < k){
                ans += cnt * l;
                k -= cnt;
            }else{
                ans += k * l;
                break;
            }
        }
        printf("%I64d\n",ans);
    }
    return 0;
}