ural 1982. Electrification Plan (最小生成树）

#include<cstdio>
#include<iostream>
#include<cstring>
#define INF 0x3f3f3f3f
using namespace std;

int edge[110][110];
int vis[110],dis[110];
bool flag[110];
int n,k,ans;

void prim()
{
int u=1,minw;
for(int i=1;i<=n;i++)
{
vis[i]=0;
dis[i]=edge[u][i];
}
vis[u]=1;
for(int i=1;i<n;i++)
{
minw=INF;
for(int j=1;j<=n;j++)
{
if(!vis[j] && dis[j]<minw)
{
minw=dis[j];
u=j;
}
}
ans+=minw;
vis[u]=1;
for(int j=1;j<=n;j++)
{
if(!vis[j] && edge[u][j]<dis[j])
dis[j]=edge[u][j];
}
}
}

int main()
{
int d;
while(scanf("%d%d",&n,&k)!=EOF)
{
memset(vis,0,sizeof(vis));
memset(flag,false,sizeof(flag));
ans=0;
for(int i=0;i<k;i++)
{
scanf("%d",&d);
flag[d]=true;
}
for(int i=1;i<=n;i++)
{
for(int j=1;j<=n;j++)
{
scanf("%d",&edge[i][j]);
}
}
for(int i=1;i<=n;i++)
{
for(int j=1;j<=n;j++)
{
if(flag[i]&&flag[j])
edge[i][j]=0;
}
}
prim();
printf("%d\n",ans);
}
return 0;
}