PAT T1016 Uniqueness of MST

dfs判断连通块的数量，prim算法建立最小生成树并判断是否唯一~

#include<bits/stdc++.h>
using namespace std;
const int maxn=1014;
const int inf=1e9;
int g[maxn][maxn];
int d[maxn];
int visit[maxn];
int N,M;
struct edge {
    int v1,v2;
}; 
vector<edge> vi;
void dfs (int s) {
    visit[s]=true;
    for (int i=1;i<=N;i++) 
    if (g[s][i]!=inf&&visit[i]==false) dfs(i);
}
int dfsTrave () {
    int block=0;
    for (int i=1;i<=N;i++) {
        if (visit[i]==false) {
            dfs(i);
            block++;
        }
    }
    return block;
}
int flag=0;
int prim (int s) {
    fill (d,d+maxn,inf);
    fill (visit,visit+maxn,0);
    d[s]=0;
    int ans=0;
    for (int i=1;i<=N;i++) {
        int u=-1,min=inf;
        for (int j=1;j<=N;j++) 
        if (visit[j]==false&&d[j]<min) {
            u=j;
            min=d[j];
        }
        if (u==-1) return -1;
        visit[u]=1;
        ans+=d[u];
        int num=0;
        for (int v=1;v<=N;v++) 
        if (visit[v]&&g[u][v]==min) num++;
        if (num>1) flag++;
        for (int v=1;v<=N;v++) {
            if (visit[v]==false&&g[u][v]!=inf&&g[u][v]<d[v]) 
            d[v]=g[u][v];
        }
    }
    return ans;
}
int main () {
    scanf ("%d %d",&N,&M);
    int u,v;
    for (int i=1;i<=N;i++) 
    for (int j=1;j<=N;j++)
    g[i][j]=inf;
    for (int i=0;i<M;i++) {
        scanf ("%d %d",&u,&v);
        scanf ("%d",&g[u][v]);
        g[v][u]=g[u][v];
    } 
    int block=dfsTrave ();
    if (block>1) {
        printf ("No MST\n");
        printf ("%d",block);
        return 0;
    }
    fill (visit,visit+maxn,false);
    int mst=prim(1);
    printf ("%d\n",mst);
    if (flag==0) printf ("Yes");
    else printf ("No");
    return 0;
}