最小全域木

Share:
 1type unionFind []int
 2
 3func newUnionFind(n int) *unionFind {
 4	uf := make(unionFind, n)
 5	for i := 0; i < n; i++ {
 6		uf[i] = -1
 7	}
 8	return &uf
 9}
10
11func (uf unionFind) Find(x int) int {
12	if uf[x] < 0 {
13		return x
14	}
15	uf[x] = uf.Find(uf[x])
16	return uf[x]
17}
18
19func (uf unionFind) Unite(x, y int) bool {
20	x, y = uf.Find(x), uf.Find(y)
21	if x != y {
22		if uf[x] > uf[y] {
23			x, y = y, x
24		}
25		uf[x] += uf[y]
26		uf[y] = x
27		return true
28	}
29	return false
30}
31
32func (uf unionFind) Size(x int) int {
33	return -uf[uf.Find(x)]
34}
35
36type edge struct{ v, to, co int }
37
38func kruskal(es []edge, n int) int {
39	uf := newUnionFind(n)
40	sort.Slice(es, func(i, j int) bool { return es[i].co < es[j].co })
41	co := 0
42	for i := 0; i < len(es); i++ {
43		if uf.Unite(es[i].v, es[i].to) {
44			co += es[i].co
45		}
46	}
47	return co
48}
49