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();
}
}