Impl UndirectedGraph

sue445 · Dec 2, 2019 · 35aa4d3 · 35aa4d3
1 parent 1e6ac39
commit 35aa4d3
Show file tree

Hide file tree

Showing 2 changed files with 77 additions and 0 deletions.
diff --git a/db/undirected_graph.go b/db/undirected_graph.go
@@ -0,0 +1,48 @@
+package db
+
+import "sort"
+
+// UndirectedGraph represents undirected graph
+type UndirectedGraph struct {
+	matrix map[string]map[string]bool
+}
+
+// NewUndirectedGraph returns a new UndirectedGraph instance
+func NewUndirectedGraph() *UndirectedGraph {
+	return &UndirectedGraph{matrix: map[string]map[string]bool{}}
+}
+
+// PutSymmetric put value to symmetric matrix
+func (g *UndirectedGraph) PutSymmetric(row string, col string, value bool) {
+	g.initRow(row)
+	g.initRow(col)
+
+	g.matrix[row][col] = value
+	g.matrix[col][row] = value
+}
+
+func (g *UndirectedGraph) initRow(row string) {
+	_, ok := g.matrix[row]
+
+	if ok {
+		return
+	}
+
+	g.matrix[row] = map[string]bool{}
+}
+
+// GetRowColumns returns columns of row
+func (g *UndirectedGraph) GetRowColumns(row string) []string {
+	var columns []string
+
+	g.initRow(row)
+
+	for k, v := range g.matrix[row] {
+		if v {
+			columns = append(columns, k)
+		}
+	}
+
+	sort.Strings(columns)
+	return columns
+}
diff --git a/db/undirected_graph_test.go b/db/undirected_graph_test.go
@@ -0,0 +1,29 @@
+package db
+
+import (
+	"github.com/stretchr/testify/assert"
+	"testing"
+)
+
+func TestUndirectedGraph_PutSymmetric(t *testing.T) {
+	g := NewUndirectedGraph()
+	g.PutSymmetric("a", "b", true)
+
+	assert.Equal(t, true, g.matrix["a"]["b"])
+	assert.Equal(t, true, g.matrix["b"]["a"])
+}
+
+func TestUndirectedGraph_GetRowColumns(t *testing.T) {
+	g := NewUndirectedGraph()
+	g.PutSymmetric("a", "b", true)
+	g.PutSymmetric("a", "c", false)
+	g.PutSymmetric("e", "a", true)
+	g.PutSymmetric("d", "a", true)
+	g.PutSymmetric("c", "d", true)
+
+	got1 := g.GetRowColumns("a")
+	assert.Equal(t, []string{"b", "d", "e"}, got1)
+
+	got2 := g.GetRowColumns("unknown")
+	assert.Empty(t, got2)
+}