prim求最小生成树

Posted 2021-01-16 sugewud

一直以来只会Kruskal

prim和dijkstra很像

只不过prim维护的是最短的边，而dijkstra维护的是最短的从起点到一个点的路径

同时prim要注意当前拓展的边是没有拓展过的

可以用堆优化

#include<bits/stdc++.h>
#define REP(i, a, b) for(register int i = (a); i < (b); i++)
#define _for(i, a, b) for(register int i = (a); i <= (b); i++)
using namespace std;

const int MAXN = 5000 + 10;
const int MAXM = 2e5 + 10;

struct Edge{ int to, w, next; };
Edge e[MAXM << 1];
int head[MAXN], d[MAXN], vis[MAXN], n, m, tot;

void AddEdge(int from, int to, int w)
{
    e[tot] = Edge{to, w, head[from]};
    head[from] = tot++;
}

void read(int& x)
{
    int f = 1; x = 0; char ch = getchar();
    while(!isdigit(ch)) { if(ch == ‘-‘) f = -1; ch = getchar(); }
    while(isdigit(ch)) { x = x * 10 + ch - ‘0‘; ch = getchar(); }
    x *= f;
}

void prim()
{
    vis[1] = 1;
    _for(i, 1, n) d[i] = i == 1 ? 0 : 1e9;
    for(int i = head[1]; ~i; i = e[i].next)
    {
        int v = e[i].to;
        d[v] = min(d[v], e[i].w);
    }
    
    int ans = 0;
    REP(k, 1, n)
    {
        int mint = 1e9, id;
        _for(i, 1, n)
            if(!vis[i] && d[i] < mint)
            {
                mint = d[i];
                id = i;
            }
        
        ans += mint;
        vis[id] = 1;
        
        for(int i = head[id]; ~i; i = e[i].next)
        {
            int v = e[i].to;
            if(vis[v]) continue;
            d[v] = min(d[v], e[i].w);
        }
    }
    printf("%d
", ans);
}

int main()
{
    memset(head, -1, sizeof(head)); tot = 0;
    read(n); read(m);
    _for(i, 1, m)
    {
        int u, v, w;
        read(u); read(v); read(w);
        AddEdge(u, v, w);
        AddEdge(v, u, w);
    }
    prim();
    return 0;
}

堆优化版本

#include<bits/stdc++.h>
#define REP(i, a, b) for(register int i = (a); i < (b); i++)
#define _for(i, a, b) for(register int i = (a); i <= (b); i++)
using namespace std;

const int MAXN = 5000 + 10;
const int MAXM = 2e5 + 10;

struct Edge{ int to, w, next; };
Edge e[MAXM << 1];
int head[MAXN], d[MAXN], vis[MAXN], n, m, tot;

void AddEdge(int from, int to, int w)
{
    e[tot] = Edge{to, w, head[from]};
    head[from] = tot++;
}

void read(int& x)
{
    int f = 1; x = 0; char ch = getchar();
    while(!isdigit(ch)) { if(ch == ‘-‘) f = -1; ch = getchar(); }
    while(isdigit(ch)) { x = x * 10 + ch - ‘0‘; ch = getchar(); }
    x *= f;
}

struct node
{
    int id, w;
    bool operator < (const node& rhs) const
    {
        return w > rhs.w;
    }
};
priority_queue<node> q;

void prim()
{
    vis[1] = 1;
    _for(i, 1, n) d[i] = i == 1 ? 0 : 1e9;
    for(int i = head[1]; ~i; i = e[i].next)
    {
        int v = e[i].to;
        d[v] = min(d[v], e[i].w);
    }
    _for(i, 1, n) q.push(node{i, d[i]});
    
    int ans = 0;
    REP(k, 1, n)
    {
        int mint, id;
        while(1)
        {
            node x = q.top(); q.pop();
            if(vis[x.id]) continue;
            id = x.id, mint = x.w;
            break;
        }

        ans += mint;
        vis[id] = 1;
        
        for(int i = head[id]; ~i; i = e[i].next)
        {
            int v = e[i].to;
            if(vis[v]) continue;
            if(d[v] > e[i].w) 
            {
                d[v] = e[i].w;
                q.push(node{v, d[v]});
            }
        }
    }
    printf("%d
", ans);
}

int main()
{
    memset(head, -1, sizeof(head)); tot = 0;
    read(n); read(m);
    _for(i, 1, m)
    {
        int u, v, w;
        read(u); read(v); read(w);
        AddEdge(u, v, w);
        AddEdge(v, u, w);
    }
    prim();
    return 0;
}