Repository: dominikbraun/graph
Pull request #101 (2023-04-17 06:51:29)

Implement spanning tree algorithms


Patch:
diff --git a/sets.go b/sets.go
index 5d0dfef..0c4aead 100644
--- a/sets.go
+++ b/sets.go
@@ -59,6 +59,61 @@ func Union[K comparable, T any](g, h Graph[K, T]) (Graph[K, T], error) {
 	return union, nil
 }
 
+// unionFind implements a union-find or disjoint set data structure that works
+// with vertex hashes as vertices. It's an internal helper type at the moment,
+// but could perhaps be exposed publicly in the future.
+//
+// unionFind is not related to the Union function.
+type unionFind[K comparable] struct {
+	parents map[K]K
+}
+
+func newUnionFind[K comparable](vertices ...K) *unionFind[K] {
+	u := &unionFind[K]{
+		parents: make(map[K]K),
+	}
+
+	for _, vertex := range vertices {
+		u.parents[vertex] = vertex
+	}
+
+	return u
+}
+
+func (u *unionFind[K]) add(vertex K) {
+	u.parents[vertex] = vertex
+}
+
+func (u *unionFind[K]) union(vertex1, vertex2 K) {
+	root1 := u.find(vertex1)
+	root2 := u.find(vertex2)
+
+	if root1 == root2 {
+		return
+	}
+
+	u.parents[root2] = root1
+}
+
+func (u *unionFind[K]) find(vertex K) K {
+	root := vertex
+
+	for u.parents[root] != root {
+		root = u.parents[root]
+	}
+
+	// Perform a path compression in order to optimize of future find calls.
+	current := vertex
+
+	for u.parents[current] != root {
+		parent := u.parents[vertex]
+		u.parents[vertex] = root
+		current = parent
+	}
+
+	return root
+}
+
 func copyVertexProperties(source VertexProperties) func(*VertexProperties) {
 	return func(p *VertexProperties) {
 		for k, v := range source.Attributes {
diff --git a/trees.go b/trees.go
new file mode 100644
index 0000000..6afe889
--- /dev/null
+++ b/trees.go
@@ -0,0 +1,83 @@
+package graph
+
+import (
+	"errors"
+	"fmt"
+	"sort"
+)
+
+// MinimumSpanningTree returns a minimum spanning tree within the given graph.
+//
+// The MST contains all vertices from the given graph as well as the required
+// edges for building the MST. The original graph remains unchanged.
+func MinimumSpanningTree[K comparable, T any](g Graph[K, T]) (Graph[K, T], error) {
+	return spanningTree(g, false)
+}
+
+// MaximumSpanningTree returns a minimum spanning tree within the given graph.
+//
+// The MST contains all vertices from the given graph as well as the required
+// edges for building the MST. The original graph remains unchanged.
+func MaximumSpanningTree[K comparable, T any](g Graph[K, T]) (Graph[K, T], error) {
+	return spanningTree(g, true)
+}
+
+func spanningTree[K comparable, T any](g Graph[K, T], maximum bool) (Graph[K, T], error) {
+	if g.Traits().IsDirected {
+		return nil, errors.New("spanning trees can only be determined for undirected graphs")
+	}
+
+	adjacencyMap, err := g.AdjacencyMap()
+	if err != nil {
+		return nil, fmt.Errorf("failed to get adjacency map: %w", err)
+	}
+
+	edges := make([]Edge[K], 0)
+	subtrees := newUnionFind[K]()
+
+	mst := New(g.(*undirected[K, T]).hash)
+
+	for v, adjacencies := range adjacencyMap {
+		vertex, properties, err := g.VertexWithProperties(v)
+		if err != nil {
+			return nil, fmt.Errorf("failed to get vertex %v: %w", v, err)
+		}
+
+		err = mst.AddVertex(vertex, copyVertexProperties(properties))
+		if err != nil {
+			return nil, fmt.Errorf("failed to add vertex %v: %w", v, err)
+		}
+
+		subtrees.add(v)
+
+		for _, edge := range adjacencies {
+			edges = append(edges, edge)
+		}
+	}
+
+	if maximum {
+		sort.Slice(edges, func(i, j int) bool {
+			return edges[i].Properties.Weight > edges[j].Properties.Weight
+		})
+	} else {
+		sort.Slice(edges, func(i, j int) bool {
+			return edges[i].Properties.Weight < edges[j].Properties.Weight
+		})
+	}
+
+	for _, edge := range edges {
+		sourceRoot := subtrees.find(edge.Source)
+		targetRoot := subtrees.find(edge.Target)
+
+		if sourceRoot != targetRoot {
+			subtrees.union(sourceRoot, targetRoot)
+
+			err := mst.AddEdge(edge.Source, edge.Target, copyEdgeProperties(edge.Properties))
+			if err != nil {
+				return nil, fmt.Errorf("failed to add edge (%v, %v): %w", edge.Source, edge.Target, err)
+			}
+		}
+	}
+
+	return mst, nil
+}


Test patch:
diff --git a/sets_test.go b/sets_test.go
index 837c4b7..8107ca4 100644
--- a/sets_test.go
+++ b/sets_test.go
@@ -198,3 +198,123 @@ func TestDirectedUnion(t *testing.T) {
 		}
 	}
 }
+
+func TestUnionFind_add(t *testing.T) {
+	tests := map[string]struct {
+		vertex         int
+		expectedParent int
+	}{
+		"add vertex 1": {
+			vertex:         1,
+			expectedParent: 1,
+		},
+	}
+
+	for name, test := range tests {
+		t.Run(name, func(t *testing.T) {
+			u := newUnionFind[int]()
+			u.add(test.vertex)
+
+			if u.parents[test.vertex] != test.expectedParent {
+				t.Errorf("expected parent %v, got %v", test.expectedParent, u.parents[test.vertex])
+			}
+		})
+	}
+}
+
+func TestUnionFind_union(t *testing.T) {
+	tests := map[string]struct {
+		vertices        []int
+		args            [2]int
+		expectedParents map[int]int
+	}{
+		"merge 1 and 2": {
+			vertices: []int{1, 2, 3, 4},
+			args:     [2]int{1, 2},
+			expectedParents: map[int]int{
+				1: 1,
+				2: 1,
+				3: 3,
+				4: 4,
+			},
+		},
+	}
+
+	for name, test := range tests {
+		t.Run(name, func(t *testing.T) {
+			u := newUnionFind(test.vertices...)
+			u.union(test.args[0], test.args[1])
+
+			for expectedKey, expectedValue := range test.expectedParents {
+				actualValue, ok := u.parents[expectedKey]
+				if !ok {
+					t.Fatalf("expected key %v not found", expectedKey)
+				}
+				if actualValue != expectedValue {
+					t.Fatalf("expected value %v for %v, got %v", expectedValue, expectedKey, actualValue)
+				}
+			}
+		})
+	}
+}
+
+func TestUnionFind_find(t *testing.T) {
+	tests := map[string]struct {
+		parents  map[int]int
+		arg      int
+		expected int
+	}{
+		"find 1 in sets 1, 2, 3": {
+			parents: map[int]int{
+				1: 1,
+				2: 2,
+				3: 3,
+			},
+			arg:      1,
+			expected: 1,
+		},
+		"find 3 in sets 1-2, 3-4": {
+			parents: map[int]int{
+				1: 1,
+				2: 1,
+				3: 3,
+				4: 3,
+			},
+			arg:      3,
+			expected: 3,
+		},
+		"find 4 in sets 1-2, 3-4": {
+			parents: map[int]int{
+				1: 1,
+				2: 1,
+				3: 3,
+				4: 3,
+			},
+			arg:      4,
+			expected: 3,
+		},
+		"find 3 in set 1-2-3": {
+			parents: map[int]int{
+				1: 1,
+				2: 1,
+				3: 2,
+			},
+			arg:      3,
+			expected: 1,
+		},
+	}
+
+	for name, test := range tests {
+		t.Run(name, func(t *testing.T) {
+			u := unionFind[int]{
+				parents: test.parents,
+			}
+
+			actual := u.find(test.arg)
+
+			if actual != test.expected {
+				t.Errorf("expected %v, got %v", test.expected, actual)
+			}
+		})
+	}
+}
diff --git a/trees_test.go b/trees_test.go
new file mode 100644
index 0000000..e206446
--- /dev/null
+++ b/trees_test.go
@@ -0,0 +1,239 @@
+package graph
+
+import (
+	"testing"
+)
+
+func TestDirectedMinimumSpanningTree(t *testing.T) {
+	tests := map[string]struct {
+		shouldFail bool
+	}{
+		"returns error": {
+			shouldFail: true,
+		},
+	}
+
+	for name, test := range tests {
+		t.Run(name, func(t *testing.T) {
+			graph := New(IntHash, Directed())
+
+			_, err := MinimumSpanningTree(graph)
+
+			if test.shouldFail != (err != nil) {
+				t.Errorf("expected error == %v, got %v", test.shouldFail, err)
+			}
+		})
+	}
+}
+
+func TestUndirectedMinimumSpanningTree(t *testing.T) {
+	tests := map[string]struct {
+		vertices                []string
+		edges                   []Edge[string]
+		expectedErr             error
+		expectedMSTAdjacencyMap map[string]map[string]Edge[string]
+	}{
+		"graph from img/mst.svg": {
+			vertices: []string{"A", "B", "C", "D"},
+			edges: []Edge[string]{
+				{Source: "A", Target: "B", Properties: EdgeProperties{Weight: 2}},
+				{Source: "A", Target: "C", Properties: EdgeProperties{Weight: 4}},
+				{Source: "A", Target: "D", Properties: EdgeProperties{Weight: 3}},
+				{Source: "B", Target: "C", Properties: EdgeProperties{Weight: 4}},
+				{Source: "B", Target: "D", Properties: EdgeProperties{Weight: 1}},
+				{Source: "C", Target: "D", Properties: EdgeProperties{Weight: 3}},
+			},
+			expectedErr: nil,
+			expectedMSTAdjacencyMap: map[string]map[string]Edge[string]{
+				"A": {
+					"B": {Source: "A", Target: "B", Properties: EdgeProperties{Weight: 2}},
+				},
+				"B": {
+					"D": {Source: "B", Target: "D", Properties: EdgeProperties{Weight: 1}},
+					"A": {Source: "B", Target: "A", Properties: EdgeProperties{Weight: 2}},
+				},
+				"C": {
+					"D": {Source: "C", Target: "D", Properties: EdgeProperties{Weight: 3}},
+				},
+				"D": {
+					"B": {Source: "D", Target: "B", Properties: EdgeProperties{Weight: 1}},
+					"C": {Source: "D", Target: "C", Properties: EdgeProperties{Weight: 3}},
+				},
+			},
+		},
+	}
+
+	for name, test := range tests {
+		t.Run(name, func(t *testing.T) {
+			g := New(StringHash)
+
+			for _, vertex := range test.vertices {
+				_ = g.AddVertex(vertex)
+			}
+
+			for _, edge := range test.edges {
+				_ = g.AddEdge(edge.Source, edge.Target, copyEdgeProperties(edge.Properties))
+			}
+
+			mst, _ := MinimumSpanningTree(g)
+			adjacencyMap, _ := mst.AdjacencyMap()
+
+			for expectedHash, expectedAdjacencies := range test.expectedMSTAdjacencyMap {
+				actualAdjacencies, ok := adjacencyMap[expectedHash]
+				if !ok {
+					t.Errorf("%s: key %v doesn't exist in adjacency map", name, expectedHash)
+					continue
+				}
+
+				for expectedAdjacency, expectedEdge := range expectedAdjacencies {
+					actualEdge, ok := actualAdjacencies[expectedAdjacency]
+					if !ok {
+						t.Errorf("%s: key %v doesn't exist in adjacencies of %v", name, expectedAdjacency, expectedHash)
+						continue
+					}
+
+					if !mst.(*undirected[string, string]).edgesAreEqual(expectedEdge, actualEdge) {
+						t.Errorf("%s: expected edge %v, got %v at AdjacencyMap[%v][%v]", name, expectedEdge, actualEdge, expectedHash, expectedAdjacency)
+					}
+
+					for expectedKey, expectedValue := range expectedEdge.Properties.Attributes {
+						actualValue, ok := actualEdge.Properties.Attributes[expectedKey]
+						if !ok {
+							t.Errorf("%s: expected attribute %v to exist in edge %v", name, expectedKey, actualEdge)
+						}
+						if actualValue != expectedValue {
+							t.Errorf("%s: expected value %v for key %v in edge %v, got %v", name, expectedValue, expectedKey, expectedEdge, actualValue)
+						}
+					}
+
+					if actualEdge.Properties.Weight != expectedEdge.Properties.Weight {
+						t.Errorf("%s: expected weight %v for edge %v, got %v", name, expectedEdge.Properties.Weight, expectedEdge, actualEdge.Properties.Weight)
+					}
+				}
+			}
+
+			for actualHash := range adjacencyMap {
+				if _, ok := test.expectedMSTAdjacencyMap[actualHash]; !ok {
+					t.Errorf("%s: unexpected key %v in union adjacency map", name, actualHash)
+				}
+			}
+		})
+	}
+}
+
+func TestDirectedMaximumSpanningTree(t *testing.T) {
+	tests := map[string]struct {
+		shouldFail bool
+	}{
+		"returns error": {
+			shouldFail: true,
+		},
+	}
+
+	for name, test := range tests {
+		t.Run(name, func(t *testing.T) {
+			graph := New(IntHash, Directed())
+
+			_, err := MaximumSpanningTree(graph)
+
+			if test.shouldFail != (err != nil) {
+				t.Errorf("expected error == %v, got %v", test.shouldFail, err)
+			}
+		})
+	}
+}
+
+func TestUndirectedMaximumSpanningTree(t *testing.T) {
+	tests := map[string]struct {
+		vertices                []string
+		edges                   []Edge[string]
+		expectedErr             error
+		expectedMSTAdjacencyMap map[string]map[string]Edge[string]
+	}{
+		"graph from img/mst.svg with higher weights": {
+			vertices: []string{"A", "B", "C", "D"},
+			edges: []Edge[string]{
+				{Source: "A", Target: "B", Properties: EdgeProperties{Weight: 20}},
+				{Source: "A", Target: "C", Properties: EdgeProperties{Weight: 4}},
+				{Source: "A", Target: "D", Properties: EdgeProperties{Weight: 3}},
+				{Source: "B", Target: "C", Properties: EdgeProperties{Weight: 4}},
+				{Source: "B", Target: "D", Properties: EdgeProperties{Weight: 10}},
+				{Source: "C", Target: "D", Properties: EdgeProperties{Weight: 30}},
+			},
+			expectedErr: nil,
+			expectedMSTAdjacencyMap: map[string]map[string]Edge[string]{
+				"A": {
+					"B": {Source: "A", Target: "B", Properties: EdgeProperties{Weight: 20}},
+				},
+				"B": {
+					"D": {Source: "B", Target: "D", Properties: EdgeProperties{Weight: 10}},
+					"A": {Source: "B", Target: "A", Properties: EdgeProperties{Weight: 20}},
+				},
+				"C": {
+					"D": {Source: "C", Target: "D", Properties: EdgeProperties{Weight: 30}},
+				},
+				"D": {
+					"B": {Source: "D", Target: "B", Properties: EdgeProperties{Weight: 10}},
+					"C": {Source: "D", Target: "C", Properties: EdgeProperties{Weight: 30}},
+				},
+			},
+		},
+	}
+
+	for name, test := range tests {
+		t.Run(name, func(t *testing.T) {
+			g := New(StringHash)
+
+			for _, vertex := range test.vertices {
+				_ = g.AddVertex(vertex)
+			}
+
+			for _, edge := range test.edges {
+				_ = g.AddEdge(edge.Source, edge.Target, copyEdgeProperties(edge.Properties))
+			}
+
+			mst, _ := MaximumSpanningTree(g)
+			adjacencyMap, _ := mst.AdjacencyMap()
+
+			for expectedHash, expectedAdjacencies := range test.expectedMSTAdjacencyMap {
+				actualAdjacencies, ok := adjacencyMap[expectedHash]
+				if !ok {
+					t.Errorf("%s: key %v doesn't exist in adjacency map", name, expectedHash)
+					continue
+				}
+
+				for expectedAdjacency, expectedEdge := range expectedAdjacencies {
+					actualEdge, ok := actualAdjacencies[expectedAdjacency]
+					if !ok {
+						t.Errorf("%s: key %v doesn't exist in adjacencies of %v", name, expectedAdjacency, expectedHash)
+						continue
+					}
+
+					if !mst.(*undirected[string, string]).edgesAreEqual(expectedEdge, actualEdge) {
+						t.Errorf("%s: expected edge %v, got %v at AdjacencyMap[%v][%v]", name, expectedEdge, actualEdge, expectedHash, expectedAdjacency)
+					}
+
+					for expectedKey, expectedValue := range expectedEdge.Properties.Attributes {
+						actualValue, ok := actualEdge.Properties.Attributes[expectedKey]
+						if !ok {
+							t.Errorf("%s: expected attribute %v to exist in edge %v", name, expectedKey, actualEdge)
+						}
+						if actualValue != expectedValue {
+							t.Errorf("%s: expected value %v for key %v in edge %v, got %v", name, expectedValue, expectedKey, expectedEdge, actualValue)
+						}
+					}
+
+					if actualEdge.Properties.Weight != expectedEdge.Properties.Weight {
+						t.Errorf("%s: expected weight %v for edge %v, got %v", name, expectedEdge.Properties.Weight, expectedEdge, actualEdge.Properties.Weight)
+					}
+				}
+			}
+
+			for actualHash := range adjacencyMap {
+				if _, ok := test.expectedMSTAdjacencyMap[actualHash]; !ok {
+					t.Errorf("%s: unexpected key %v in union adjacency map", name, actualHash)
+				}
+			}
+		})
+	}
+}