HDU 4313 Matrix 类似 克鲁斯卡尔算法

#include <iostream>
#include <cstdio>
#include <cstdlib>
#include <cmath>
#include <cstring>
#include <string>
#include <algorithm>
#include <set>
#include<map>
#include<ctime>
using namespace std;
set<int> Set;
struct Edge
{
    int x,y,val;
    bool operator<(const Edge & a)const
    {
        return val>a.val;
    }
}edge[100009];
int n,k,fa[100009];
int find(int x)
{
    return x==fa[x]?x:fa[x]=find(fa[x]);
}
void solve()
{
    int x,y,now=0;
    long long ans=0;
    for(int i=0;i<=n;i++)
    {
        fa[i]=i;
    }
    for(int i=0;i<n-1;i++)
    {
        x=edge[i].x;
        y=edge[i].y;
        x=find(x);
        y=find(y);
        if(Set.count(x)!=0&&Set.count(y)!=0)
        {
            now++;
            ans+=edge[i].val;
            if(now==k-1)break;
            continue;
        }
        if(Set.count(x)!=0)
        {
            fa[y]=x;
        }
        else
        {
            fa[x]=y;
        }
    }
    printf("%I64d\n",ans);
}
int main()
{
    int t,f;
    scanf("%d",&t);
    while(t--)
    {
        scanf("%d%d",&n,&k);
        Set.clear();
        for(int i=0;i<n-1;i++)
        {
            scanf("%d%d%d",&edge[i].x,&edge[i].y,&edge[i].val);
        }
        sort(edge,edge+n-1);
        for(int i=0;i<k;i++)
        {
            scanf("%d",&f);
            Set.insert(f);
        }
        solve();
    }
}