Golang算法模板-最小生成树模板Kruskal/Prim

154 阅读2分钟

Kruskal

洛谷 P3366【模板】最小生成树

https://www.luogu.com.cn/problem/P3366

Kruskal算法优先试用于稀疏图 把n个点中m条边的边权按照从小到大排序,然后遍历这m条边,如果这一条边两端的两个不连通,就合并,如果连通了,这条边就作废(这里用到了并查集的知识)。一直到用到的边数为n-1条,因为生成一个树只需要n-1条边。

package main
​
import (
    "fmt"
    "sort"
)
​
type Edge struct {
    u, v, w int
}
​
var edge = [200005]Edge{}
var pre = [5005]int{}
​
func find(x int) int {
    r := x
    for pre[r] != r {
        r = pre[r]
    }
    for x != r {
        next := pre[x]
        pre[x] = r
        x = next
    }
    return r
}
func join(a, b int) {
    fa := find(a)
    fb := find(b)
    if fa != fb {
        pre[fa] = fb
    }
}
​
func main() {
    sum, count := 0, 0
    var n, m int
    fmt.Scan(&n, &m)
    for i := 1; i <= n; i++ {
        pre[i] = i
    }
    for i := 1; i <= m; i++ {
        fmt.Scan(&edge[i].u, &edge[i].v, &edge[i].w)
    }
    sort.Slice(edge[:m+1], func(i, j int) bool {    //注意这里要从0开始
        return edge[i].w < edge[j].w
    })
    //fmt.Println(edge[:20])
​
    for i := 1; i <= m; i++ {
        if find(edge[i].u) != find(edge[i].v) {
            join(edge[i].u, edge[i].v)
            sum += edge[i].w
            count++
            if count == n-1 {
                break
            }
        }
    }
    if count == n-1 {
        fmt.Printf("%d", sum)
    } else {
        fmt.Printf("orz")
    }
}

Prim

Prim算法,(和最短路径DJ算法很像),优先试用于稠密图 一开始n个点相互独立,如果要让n个点相互连通,那么让这n个点都依次进入树里面。 从任意一点开始构造生成树,这里先从1点开始,让1点入树,book标记,dis标记到树的最短距离,然后遍历n个点,找到到树的最短距离,这个点入树,标记,然后更新dis数组,直到所有点都入树。

package mainimport (
    "fmt"
    "math"
)
​
var edge [5005][5005]int
var book [5005]int
var dis [5005]int
var inf = math.MaxIntfunc main() {
    //fmt.Println(inf)
    count, sum := 0, 0
    var n, m int
    fmt.Scan(&n, &m)
    for i := 1; i <= n; i++ { //初始化图
        for j := 1; j <= n; j++ {
            if i == j {
                edge[i][j] = 0
            } else {
                edge[i][j] = inf
            }
        }
    }
    for i := 1; i <= m; i++ {
        var u, v, w int
        fmt.Scan(&u, &v, &w)
        if w < edge[u][v] {
            edge[u][v] = w
            edge[v][u] = w
        }
    }
    for j := 1; j <= n; j++ { //先让1号进树
        dis[j] = edge[1][j]
    }
    book[1] = 1
    count++
​
    for r := 1; r <= n-1; r++ { //还需要再入n-1个点,循环n-1次
        var nextPos int = -1
        var minn = inf
        for i := 1; i <= n; i++ {
            if book[i] == 0 && dis[i] < minn {
                minn = dis[i]
                nextPos = i
            }
        }
        if nextPos == -1 {
            break
        }
        book[nextPos] = 1
        count++
        sum += dis[nextPos]
        for k := 1; k <= n; k++ {
            if book[k] == 0 && dis[k] > edge[nextPos][k] {
                dis[k] = edge[nextPos][k]
            }
        }
    }
    if count == n {
        fmt.Printf("%d", sum)
    } else {
        fmt.Printf("orz")
    }
}
​