From 0254892642fa4a9c2e55bec0a152bb721830b455 Mon Sep 17 00:00:00 2001 From: jafar75 Date: Fri, 15 Mar 2024 20:55:06 +0330 Subject: [PATCH] feat: add prim algorithm to find mst (#710) --- README.md | 13 +++-- graph/prim.go | 58 ++++++++++++++++++ graph/prim_test.go | 143 +++++++++++++++++++++++++++++++++++++++++++++ 3 files changed, 208 insertions(+), 6 deletions(-) create mode 100644 graph/prim.go create mode 100644 graph/prim_test.go diff --git a/README.md b/README.md index 1d9b007cc..a7ef32fb5 100644 --- a/README.md +++ b/README.md @@ -455,12 +455,13 @@ Read our [Contribution Guidelines](CONTRIBUTING.md) before you contribute. 5. [`FloydWarshall`](./graph/floydwarshall.go#L15): FloydWarshall Returns all pair's shortest path using Floyd Warshall algorithm 6. [`GetIdx`](./graph/depthfirstsearch.go#L3): No description provided. 7. [`KruskalMST`](./graph/kruskal.go#L23): No description provided. -8. [`LowestCommonAncestor`](./graph/lowestcommonancestor.go#L111): For each node, we will precompute its ancestor above him, its ancestor two nodes above, its ancestor four nodes above, etc. Let's call `jump[j][u]` is the `2^j`-th ancestor above the node `u` with `u` in range `[0, numbersVertex)`, `j` in range `[0,MAXLOG)`. These information allow us to jump from any node to any ancestor above it in `O(MAXLOG)` time. -9. [`New`](./graph/graph.go#L16): Constructor functions for graphs (undirected by default) -10. [`NewTree`](./graph/lowestcommonancestor.go#L84): No description provided. -11. [`NewUnionFind`](./graph/unionfind.go#L24): Initialise a new union find data structure with s nodes -12. [`NotExist`](./graph/depthfirstsearch.go#L12): No description provided. -13. [`Topological`](./graph/topological.go#L7): Topological assumes that graph given is valid and that its possible to get a topological ordering. constraints are array of []int{a, b}, representing an edge going from a to b +8. [`PrimMST`](./graph/prim.go#30): Computes the minimum spanning tree of a weighted undirected graph +9. [`LowestCommonAncestor`](./graph/lowestcommonancestor.go#L111): For each node, we will precompute its ancestor above him, its ancestor two nodes above, its ancestor four nodes above, etc. Let's call `jump[j][u]` is the `2^j`-th ancestor above the node `u` with `u` in range `[0, numbersVertex)`, `j` in range `[0,MAXLOG)`. These information allow us to jump from any node to any ancestor above it in `O(MAXLOG)` time. +10. [`New`](./graph/graph.go#L16): Constructor functions for graphs (undirected by default) +11. [`NewTree`](./graph/lowestcommonancestor.go#L84): No description provided. +12. [`NewUnionFind`](./graph/unionfind.go#L24): Initialise a new union find data structure with s nodes +13. [`NotExist`](./graph/depthfirstsearch.go#L12): No description provided. +14. [`Topological`](./graph/topological.go#L7): Topological assumes that graph given is valid and that its possible to get a topological ordering. constraints are array of []int{a, b}, representing an edge going from a to b --- ##### Types diff --git a/graph/prim.go b/graph/prim.go new file mode 100644 index 000000000..d99dcd6e1 --- /dev/null +++ b/graph/prim.go @@ -0,0 +1,58 @@ +// The Prim's algorithm computes the minimum spanning tree for a weighted undirected graph +// Worst Case Time Complexity: O(E log V) using Binary heap, where V is the number of vertices and E is the number of edges +// Space Complexity: O(V + E) +// Implementation is based on the book 'Introduction to Algorithms' (CLRS) + +package graph + +import ( + "container/heap" +) + +type minEdge []Edge + +func (h minEdge) Len() int { return len(h) } +func (h minEdge) Less(i, j int) bool { return h[i].Weight < h[j].Weight } +func (h minEdge) Swap(i, j int) { h[i], h[j] = h[j], h[i] } + +func (h *minEdge) Push(x interface{}) { + *h = append(*h, x.(Edge)) +} + +func (h *minEdge) Pop() interface{} { + old := *h + n := len(old) + x := old[n-1] + *h = old[0 : n-1] + return x +} + +func (g *Graph) PrimMST(start Vertex) ([]Edge, int) { + var mst []Edge + marked := make([]bool, g.vertices) + h := &minEdge{} + // Pushing neighbors of the start node to the binary heap + for neighbor, weight := range g.edges[int(start)] { + heap.Push(h, Edge{start, Vertex(neighbor), weight}) + } + marked[start] = true + cost := 0 + for h.Len() > 0 { + e := heap.Pop(h).(Edge) + end := int(e.End) + // To avoid cycles + if marked[end] { + continue + } + marked[end] = true + cost += e.Weight + mst = append(mst, e) + // Check for neighbors of the newly added edge's End vertex + for neighbor, weight := range g.edges[end] { + if !marked[neighbor] { + heap.Push(h, Edge{e.End, Vertex(neighbor), weight}) + } + } + } + return mst, cost +} diff --git a/graph/prim_test.go b/graph/prim_test.go new file mode 100644 index 000000000..8a3d42c12 --- /dev/null +++ b/graph/prim_test.go @@ -0,0 +1,143 @@ +package graph + +import ( + "fmt" + "reflect" + "testing" +) + +func TestPrimMST(t *testing.T) { + + var testCases = []struct { + edges []Edge + vertices int + start int + cost int + mst []Edge + }{ + { + edges: []Edge{ + {Start: 0, End: 1, Weight: 4}, + {Start: 0, End: 2, Weight: 13}, + {Start: 0, End: 3, Weight: 7}, + {Start: 0, End: 4, Weight: 7}, + {Start: 1, End: 2, Weight: 9}, + {Start: 1, End: 3, Weight: 3}, + {Start: 1, End: 4, Weight: 7}, + {Start: 2, End: 3, Weight: 10}, + {Start: 2, End: 4, Weight: 14}, + {Start: 3, End: 4, Weight: 4}, + }, + vertices: 5, + start: 0, + cost: 20, + mst: []Edge{ + {Start: 0, End: 1, Weight: 4}, + {Start: 1, End: 3, Weight: 3}, + {Start: 3, End: 4, Weight: 4}, + {Start: 1, End: 2, Weight: 9}, + }, + }, + { + edges: []Edge{ + {Start: 0, End: 1, Weight: 4}, + {Start: 0, End: 7, Weight: 8}, + {Start: 1, End: 2, Weight: 8}, + {Start: 1, End: 7, Weight: 11}, + {Start: 2, End: 3, Weight: 7}, + {Start: 2, End: 5, Weight: 4}, + {Start: 2, End: 8, Weight: 2}, + {Start: 3, End: 4, Weight: 9}, + {Start: 3, End: 5, Weight: 14}, + {Start: 4, End: 5, Weight: 10}, + {Start: 5, End: 6, Weight: 2}, + {Start: 6, End: 7, Weight: 1}, + {Start: 6, End: 8, Weight: 6}, + {Start: 7, End: 8, Weight: 7}, + }, + vertices: 9, + start: 3, + cost: 37, + mst: []Edge{ + {Start: 3, End: 2, Weight: 7}, + {Start: 2, End: 8, Weight: 2}, + {Start: 2, End: 5, Weight: 4}, + {Start: 5, End: 6, Weight: 2}, + {Start: 6, End: 7, Weight: 1}, + {Start: 2, End: 1, Weight: 8}, + {Start: 1, End: 0, Weight: 4}, + {Start: 3, End: 4, Weight: 9}, + }, + }, + { + edges: []Edge{ + {Start: 0, End: 1, Weight: 2}, + {Start: 0, End: 3, Weight: 6}, + {Start: 1, End: 2, Weight: 3}, + {Start: 1, End: 3, Weight: 8}, + {Start: 1, End: 4, Weight: 5}, + {Start: 2, End: 4, Weight: 7}, + {Start: 3, End: 4, Weight: 9}, + }, + vertices: 5, + start: 2, + cost: 16, + mst: []Edge{ + {Start: 2, End: 1, Weight: 3}, + {Start: 1, End: 0, Weight: 2}, + {Start: 1, End: 4, Weight: 5}, + {Start: 0, End: 3, Weight: 6}, + }, + }, + { + edges: []Edge{ + {Start: 0, End: 0, Weight: 0}, + }, + vertices: 1, + start: 0, + cost: 0, + mst: nil, + }, + { + edges: []Edge{ + {Start: 0, End: 1, Weight: 1}, + {Start: 0, End: 2, Weight: 6}, + {Start: 0, End: 3, Weight: 5}, + {Start: 1, End: 2, Weight: 2}, + {Start: 1, End: 4, Weight: 4}, + {Start: 2, End: 4, Weight: 9}, + }, + vertices: 5, + start: 4, + cost: 12, + mst: []Edge{ + {Start: 4, End: 1, Weight: 4}, + {Start: 1, End: 0, Weight: 1}, + {Start: 1, End: 2, Weight: 2}, + {Start: 0, End: 3, Weight: 5}, + }, + }, + } + + for i, testCase := range testCases { + t.Run(fmt.Sprintf("Test Case %d", i), func(t *testing.T) { + // Initializing graph, adding edges + graph := New(testCase.vertices) + graph.Directed = false + for _, edge := range testCase.edges { + graph.AddWeightedEdge(int(edge.Start), int(edge.End), edge.Weight) + } + + computedMST, computedCost := graph.PrimMST(Vertex(testCase.start)) + + // Compare the computed result with the expected result + if computedCost != testCase.cost { + t.Errorf("Test Case %d, Expected Cost: %d, Computed: %d", i, testCase.cost, computedCost) + } + + if !reflect.DeepEqual(testCase.mst, computedMST) { + t.Errorf("Test Case %d, Expected MST: %v, Computed: %v", i, testCase.mst, computedMST) + } + }) + } +}