次小生成树模板

ural 1416 求最小生成树和次小生成树

#include <bits/stdc++.h>
using namespace std;
#define inf 0x3f3f3f3f
#define PI acos(-1.0)

typedef long long ll;
const int maxn=505;
int mp[maxn][maxn],cost[maxn],n,m,pre[maxn],maxe[maxn][maxn];
bool vis[maxn],inmst[maxn][maxn];

int Mst()
{
    memset(vis,0,sizeof(vis));
    memset(cost,inf,sizeof(cost));\
    memset(maxe,0,sizeof(maxe));
    cost[1]=0;
    vis[1]=true;
    int sum=0;
    for(int i=1;i<=n;i++)
        cost[i]=mp[1][i],pre[i]=1;
    for(int i=1;i<n;i++)
    {
        int minn=inf;
        int k=-1;
        for(int j=1;j<=n;j++)
        {
            if(!vis[j]&&cost[j]<minn)
                minn=cost[j],k=j;
        }
        if(k==-1)
            return -1;
        vis[k]=true;
        inmst[pre[k]][k]=inmst[k][pre[k]]=1;
        sum+=minn;
        for(int j=1;j<=n;j++)
        {
            if(vis[j]&&j!=k)
                maxe[j][k]=maxe[k][j]=max(maxe[pre[k]][j],cost[k]);
            if(!vis[j]&&cost[j]>mp[k][j])
                cost[j]=mp[k][j],pre[j]=k;
        }
    }
    return sum;
}

int Smst(int x)
{
    int ans=inf;
    for(int i=1;i<=n;i++)
    {
        for(int j=1;j<i;j++)
        {
            if(mp[i][j]!=inf&&!inmst[i][j])
                ans=min(ans,x-maxe[i][j]+mp[i][j]);
        }
    }
    if(ans==inf)
        return -1;
    return ans;
}

int main()
{
    while(scanf("%d%d",&n,&m)!=EOF)
    {
        memset(mp,inf,sizeof(mp));
        memset(inmst,0,sizeof(inmst));
        for(int i=0;i<m;i++)
        {
            int a,b,c;
            cin >> a >> b >> c;
            mp[a][b]=mp[b][a]=c;
        }
        int a=Mst();
        if(a==-1)
        {
            printf("Cost: -1\nCost: -1\n");
            continue;
        }
        int b=Smst(a);
        printf("Cost: %d\nCost: %d\n",a,b);
    }
    return 0;
}