/
spanning_tree.go
108 lines (86 loc) · 2.65 KB
/
spanning_tree.go
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
// Copyright ©2014 The gonum Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.
package path
import (
"sort"
"github.com/gonum/graph"
"github.com/gonum/graph/concrete"
"github.com/gonum/graph/internal"
)
// EdgeListerGraph is an undirected graph than returns its complete set of edges.
type EdgeListerGraph interface {
graph.Undirected
Edges() []graph.Edge
}
// Prim generates a minimum spanning tree of g by greedy tree extension, placing
// the result in the destination. The destination is not cleared first.
func Prim(dst graph.MutableUndirected, g EdgeListerGraph) {
var weight graph.WeightFunc
if g, ok := g.(graph.Weighter); ok {
weight = g.Weight
} else {
weight = graph.UniformCost
}
nlist := g.Nodes()
if nlist == nil || len(nlist) == 0 {
return
}
dst.AddNode(nlist[0])
remainingNodes := make(internal.IntSet)
for _, node := range nlist[1:] {
remainingNodes.Add(node.ID())
}
edgeList := g.Edges()
for remainingNodes.Count() != 0 {
var edges []concrete.WeightedEdge
for _, edge := range edgeList {
if (dst.Has(edge.From()) && remainingNodes.Has(edge.To().ID())) ||
(dst.Has(edge.To()) && remainingNodes.Has(edge.From().ID())) {
edges = append(edges, concrete.WeightedEdge{Edge: edge, Cost: weight(edge)})
}
}
sort.Sort(byWeight(edges))
myEdge := edges[0]
dst.SetEdge(myEdge.Edge, myEdge.Cost)
remainingNodes.Remove(myEdge.Edge.From().ID())
}
}
// Kruskal generates a minimum spanning tree of g by greedy tree coalesence, placing
// the result in the destination. The destination is not cleared first.
func Kruskal(dst graph.MutableUndirected, g EdgeListerGraph) {
var weight graph.WeightFunc
if g, ok := g.(graph.Weighter); ok {
weight = g.Weight
} else {
weight = graph.UniformCost
}
edgeList := g.Edges()
edges := make([]concrete.WeightedEdge, 0, len(edgeList))
for _, edge := range edgeList {
edges = append(edges, concrete.WeightedEdge{Edge: edge, Cost: weight(edge)})
}
sort.Sort(byWeight(edges))
ds := newDisjointSet()
for _, node := range g.Nodes() {
ds.makeSet(node.ID())
}
for _, edge := range edges {
// The disjoint set doesn't really care for which is head and which is tail so this
// should work fine without checking both ways
if s1, s2 := ds.find(edge.Edge.From().ID()), ds.find(edge.Edge.To().ID()); s1 != s2 {
ds.union(s1, s2)
dst.SetEdge(edge.Edge, edge.Cost)
}
}
}
type byWeight []concrete.WeightedEdge
func (e byWeight) Len() int {
return len(e)
}
func (e byWeight) Less(i, j int) bool {
return e[i].Cost < e[j].Cost
}
func (e byWeight) Swap(i, j int) {
e[i], e[j] = e[j], e[i]
}