最小生成树----prim算法的堆优化

题目描述

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

输入输出格式

输入格式：

 

第一行包含两个整数N、M，表示该图共有N个结点和M条无向边。（N<=5000，M<=200000）

接下来M行每行包含三个整数Xi、Yi、Zi，表示有一条长度为Zi的无向边连接结点Xi、Yi

 

输出格式：

 

输出包含一个数，即最小生成树的各边的长度之和；如果该图不连通则输出orz

 

#include <cstdio>
#include <cstring>
#include <cctype>
#include <stdlib.h>
#include <string>
#include <map>
#include <iostream>
#include <set>
#include <stack>
#include <cmath>
#include <queue>
#include <vector>
#include <algorithm>
using namespace std;
#define mem(a,b) memset(a,b,sizeof(a))
typedef long long ll;
typedef pair<int,int> pir;
const int N=5000+10;
const int M=200000+10;

int first[N],tot;
int vis[N],dis[N],n,m;
priority_queue <pir,vector<pir>,greater<pir> >q;
struct edge
{
    int v,w,next;
} e[M*2];

void add_edge(int u,int v,int w)
{
    e[tot].v=v;
    e[tot].w=w;
    e[tot].next=first[u];
    first[u]=tot++;
}
void init()
{
    mem(first,-1);
    tot=0;
    mem(dis,127);
}
void prim()
{
    int cnt=0,sum=0;
    dis[1]=0;
    q.push(make_pair(0,1));
    while(!q.empty()&&cnt<n)
    {
        int d=q.top().first,u=q.top().second;
        q.pop();
        if(!vis[u])
        {
            cnt++;
            sum+=d;
            vis[u]=1;
            for(int i=first[u]; ~i; i=e[i].next)
                if(e[i].w<dis[e[i].v])
                {
                    dis[e[i].v]=e[i].w;
                    q.push(make_pair(dis[e[i].v],e[i].v));
                }
        }
    }
    if(cnt==n) printf("%d\n",sum);
    else puts("orz");
}
int main()
{
    int u,v,w;
    init();
    scanf("%d%d",&n,&m);
    for(int i=1; i<=m; i++)
    {
        scanf("%d%d%d",&u,&v,&w);
        add_edge(u,v,w);
        add_edge(v,u,w);
    }
    prim();
    return 0;
}