Trees in Golang
April 08, 2019
DFS and BFS tree implementations with Golang.
Setting up the test
Set up trees_test.go
with the following file:
package trees
import "testing"
func TestGetTreeRoot(t *testing.T) {
n1 := Node{8, []Node{}}
tree := Tree{}
tree.setRoot(&n1)
expected := &n1
if observed := tree.getRoot(); observed != expected {
t.Fatalf("Trees()) = %v, want %v", observed, expected)
}
}
func TestAppendChild(t *testing.T) {
n1 := Node{8, []Node{}}
tree := Tree{}
tree.setRoot(&n1)
n2 := Node{2, []Node{}}
n3 := Node{3, []Node{}}
n1.append(&n2)
n1.append(&n3)
expected := 2
if observed := n1.childLength(); observed != expected {
t.Fatalf("Trees()) = %v, want %v", observed, expected)
}
}
// This test only compares returned values, not addresses
func TestBFS(t *testing.T) {
n1 := Node{8, []Node{}}
tree := Tree{}
tree.setRoot(&n1)
n4 := Node{4, []Node{}}
n5 := Node{5, []Node{}}
n2 := Node{2, []Node{n4}}
n3 := Node{3, []Node{n5}}
n1.append(&n2)
n1.append(&n3)
expected := []int{n1.data, n2.data, n3.data, n4.data, n5.data}
observed := tree.bfs()
for i := range observed {
if observed[i] != expected[i] {
t.Fatalf("Failed: bfs => %d, want %d", observed[i], expected[i])
}
t.Logf("Success: bfs => %d, want %d", observed[i], expected[i])
}
}
func TestDFS(t *testing.T) {
n1 := Node{8, []Node{}}
tree := Tree{}
tree.setRoot(&n1)
n4 := Node{4, []Node{}}
n5 := Node{5, []Node{}}
n6 := Node{6, []Node{}}
n2 := Node{2, []Node{n4, n6}}
n3 := Node{3, []Node{n5}}
n1.append(&n2)
n1.append(&n3)
expected := []int{n1.data, n2.data, n4.data, n6.data, n3.data, n5.data}
observed := tree.dfs()
for i := range observed {
if observed[i] != expected[i] {
t.Fatalf("Failed: bfs => %d, want %d", observed[i], expected[i])
}
t.Logf("Success: bfs => %d, want %d", observed[i], expected[i])
}
}
Trees implementation
package trees
// Tree is a basic tree structure
type Tree struct {
root *Node
}
// Node represents a graph vertex with a data point and children
type Node struct {
data int
children []Node
}
// Trees should have a comment documenting it.
func (t *Tree) setRoot(n *Node) {
t.root = n
}
func (t *Tree) getRoot() *Node {
return t.root
}
func (n *Node) append(c *Node) {
n.children = append(n.children, *c)
}
func (n *Node) childLength() int {
return len(n.children)
}
func (t *Tree) bfs() []int {
n := t.getRoot()
if n == nil {
return []int{}
}
arr := []Node{*n}
res := []int{}
// Iterate through, if children, push to end of array
for len(arr) > 0 {
x := arr[0]
arr = arr[1:]
res = append(res, x.data)
arr = append(arr, x.children...)
}
return res
}
func (t *Tree) dfs() []int {
n := t.getRoot()
if n == nil {
return []int{}
}
arr := []Node{*n}
res := []int{}
// Iterate through, if children, push to end of array
for len(arr) > 0 {
x := arr[0]
arr = arr[1:]
res = append(res, x.data)
arr = append(x.children, arr...)
}
return res
}
Running Tests
In the directory, run go test
.
Related Articles
A personal blog on all things of interest. Written by Dennis O'Keeffe, Follow me on Twitter