gitgitcode
diff --git a/‎graph/kruskal.go‎
Lines changed: 113 additions & 0 deletions b/‎graph/kruskal.go‎
Lines changed: 113 additions & 0 deletions
diff --git a/‎graph/kruskal_test.go‎
Lines changed: 162 additions & 0 deletions b/‎graph/kruskal_test.go‎
Lines changed: 162 additions & 0 deletions
@@ -0,0 +1,113 @@
+// KRUSKAL'S ALGORITHM
+// https://cp-algorithms.com/data_structures/disjoint_set_union.html
+// https://cp-algorithms.com/graph/mst_kruskal_with_dsu.html
+
+package graph
+
+import (
+	"sort"
+)
+
+type Vertex int
+
+// Edge describes the edge of a weighted graph
+type Edge struct {
+	Start  Vertex
+	End    Vertex
+	Weight int
+}
+
+// DisjointSetUnionElement describes what an element of DSU looks like
+type DisjointSetUnionElement struct {
+	Parent Vertex
+	Rank   int
+}
+
+// DisjointSetUnion is a data structure that treats its elements as separate sets
+// and provides fast operations for set creation, merging sets, and finding the parent
+// of the given element of a set.
+type DisjointSetUnion []DisjointSetUnionElement
+
+// NewDSU will return an initialised DSU using the value of n
+// which will be treated as the number of elements out of which
+// the DSU is being made
+func NewDSU(n int) *DisjointSetUnion {
+
+	dsu := DisjointSetUnion(make([]DisjointSetUnionElement, n))
+	return &dsu
+}
+
+// MakeSet will create a set in the DSU for the given node
+func (dsu DisjointSetUnion) MakeSet(node Vertex) {
+
+	dsu[node].Parent = node
+	dsu[node].Rank = 0
+}
+
+// FindSetRepresentative will return the parent element of the set the given node
+// belongs to. Since every single element in the path from node to parent
+// has the same parent, we store the parent value for each element in the
+// path. This reduces consequent function calls and helps in going from O(n)
+// to O(log n). This is known as path compression technique.
+func (dsu DisjointSetUnion) FindSetRepresentative(node Vertex) Vertex {
+
+	if node == dsu[node].Parent {
+		return node
+	}
+
+	dsu[node].Parent = dsu.FindSetRepresentative(dsu[node].Parent)
+	return dsu[node].Parent
+}
+
+// unionSets will merge two given sets. The naive implementation of this
+// always combines the secondNode's tree with the firstNode's tree. This can lead
+// to creation of trees of length O(n) so we optimize by attaching the node with
+// smaller rank to the node with bigger rank. Rank represents the upper bound depth of the tree.
+func (dsu DisjointSetUnion) UnionSets(firstNode Vertex, secondNode Vertex) {
+
+	firstNode = dsu.FindSetRepresentative(firstNode)
+	secondNode = dsu.FindSetRepresentative(secondNode)
+
+	if firstNode != secondNode {
+
+		if dsu[firstNode].Rank < dsu[secondNode].Rank {
+			firstNode, secondNode = secondNode, firstNode
+		}
+		dsu[secondNode].Parent = firstNode
+
+		if dsu[firstNode].Rank == dsu[secondNode].Rank {
+			dsu[firstNode].Rank++
+		}
+	}
+}
+
+// KruskalMST will return a minimum spanning tree along with its total cost
+// to using Kruskal's algorithm. Time complexity is O(m * log (n)) where m is
+// the number of edges in the graph and n is number of nodes in it.
+func KruskalMST(n int, edges []Edge) ([]Edge, int) {
+
+	var mst []Edge // The resultant minimum spanning tree
+	var cost int = 0
+
+	dsu := NewDSU(n)
+
+	for i := 0; i < n; i++ {
+		dsu.MakeSet(Vertex(i))
+	}
+
+	sort.SliceStable(edges, func(i, j int) bool {
+		return edges[i].Weight < edges[j].Weight
+	})
+
+	for _, edge := range edges {
+
+		if dsu.FindSetRepresentative(edge.Start) != dsu.FindSetRepresentative(edge.End) {
+
+			mst = append(mst, edge)
+			cost += edge.Weight
+			dsu.UnionSets(edge.Start, edge.End)
+		}
+	}
+
+	return mst, cost
+}
@@ -0,0 +1,162 @@
+package graph
+
+import (
+	"fmt"
+	"testing"
+)
+
+func Test_KruskalMST(t *testing.T) {
+
+	var testCases = []struct {
+		n     int
+		graph []Edge
+		cost  int
+	}{
+		{
+			n: 5,
+			graph: []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,
+				},
+			},
+			cost: 20,
+		},
+		{
+			n: 3,
+			graph: []Edge{
+				{
+					Start:  0,
+					End:    1,
+					Weight: 12,
+				},
+				{
+					Start:  0,
+					End:    2,
+					Weight: 18,
+				},
+				{
+					Start:  1,
+					End:    2,
+					Weight: 6,
+				},
+			},
+			cost: 18,
+		},
+		{
+			n: 4,
+			graph: []Edge{
+				{
+					Start:  0,
+					End:    1,
+					Weight: 2,
+				},
+				{
+					Start:  0,
+					End:    2,
+					Weight: 1,
+				},
+				{
+					Start:  0,
+					End:    3,
+					Weight: 2,
+				},
+				{
+					Start:  1,
+					End:    2,
+					Weight: 1,
+				},
+				{
+					Start:  1,
+					End:    3,
+					Weight: 2,
+				},
+				{
+					Start:  2,
+					End:    3,
+					Weight: 3,
+				},
+			},
+			cost: 4,
+		},
+		{
+			n: 2,
+			graph: []Edge{
+				{
+					Start:  0,
+					End:    1,
+					Weight: 4000000,
+				},
+			},
+			cost: 4000000,
+		},
+		{
+			n: 1,
+			graph: []Edge{
+				{
+					Start:  0,
+					End:    0,
+					Weight: 0,
+				},
+			},
+			cost: 0,
+		},
+	}
+
+	for i := range testCases {
+
+		t.Run(fmt.Sprintf("Test Case %d", i), func(t *testing.T) {
+
+			_, computed := KruskalMST(testCases[i].n, testCases[i].graph)
+			if computed != testCases[i].cost {
+				t.Errorf("Test Case %d, Expected: %d, Computed: %d", i, testCases[i].cost, computed)
+			}
+		})
+	}
+}