prim(邻接表）P3366 【模板】最小生成树

如题，给出一个无向图，求出最小生成树，如果该图不连通，则输出 orz。

#include <bits/stdc++.h>
#define ll long long
#define inf 0x3f3f3f3f

using namespace std;
const int N=5e3+5;
struct node
{
    int v,w;
};
vector<node> a[N];
int n,m,u,v,w,vis[N],dis[N],cnt=0;

int prim(int s)
{
    vis[s]=1;
    dis[s]=0;
    for(int i=0;i<a[s].size();i++)
    {
        int v=a[s][i].v,w=a[s][i].w;
        if(vis[v]==0)
            dis[v]=min(dis[v],w);
    }
    int res=0;
    for(int i=0;i<n-1;i++)
    {
        int mi=inf,k=0;
        for(int j=1;j<=n;j++)
        {
            if(vis[j]==0&&mi>dis[j])
            {
                mi=dis[j];
                k=j;
            }
        }
        if(k==0)
            break;
        cnt++;
        vis[k]=1;
        res+=mi;
        for(int j=0;j<a[k].size();j++)
        {
            int v=a[k][j].v,w=a[k][j].w;
            if(vis[v]==0)
                dis[v]=min(dis[v],w);
        }
    }
    return res;
}
int main()
{
    cin >> n >> m;
    memset(vis,0,sizeof(vis));
    memset(dis,0x3f,sizeof(dis));
    while(m--)
    {
        cin >> u >> v >> w;
        a[u].push_back({v,w});
        a[v].push_back({u,w});
    }
    int t=prim(1);
    if(cnt==n-1)
        printf("%d\n",t);
    else
        printf("orz\n");
    return 0;
}