| 28 | } |
| 29 | |
| 30 | func (g *Graph) PrimMST(start Vertex) ([]Edge, int) { |
| 31 | var mst []Edge |
| 32 | marked := make([]bool, g.vertices) |
| 33 | h := &minEdge{} |
| 34 | // Pushing neighbors of the start node to the binary heap |
| 35 | for neighbor, weight := range g.edges[int(start)] { |
| 36 | heap.Push(h, Edge{start, Vertex(neighbor), weight}) |
| 37 | } |
| 38 | marked[start] = true |
| 39 | cost := 0 |
| 40 | for h.Len() > 0 { |
| 41 | e := heap.Pop(h).(Edge) |
| 42 | end := int(e.End) |
| 43 | // To avoid cycles |
| 44 | if marked[end] { |
| 45 | continue |
| 46 | } |
| 47 | marked[end] = true |
| 48 | cost += e.Weight |
| 49 | mst = append(mst, e) |
| 50 | // Check for neighbors of the newly added edge's End vertex |
| 51 | for neighbor, weight := range g.edges[end] { |
| 52 | if !marked[neighbor] { |
| 53 | heap.Push(h, Edge{e.End, Vertex(neighbor), weight}) |
| 54 | } |
| 55 | } |
| 56 | } |
| 57 | return mst, cost |
| 58 | } |