diff --git a/graph/traverse/bf_iterator.go b/graph/traverse/bf_iterator.go new file mode 100644 index 0000000..12932f4 --- /dev/null +++ b/graph/traverse/bf_iterator.go @@ -0,0 +1,65 @@ +package traverse + +import "github.com/hmdsefi/gograph/graph" + +type breadthFirstIterator[T comparable] struct { + graph graph.Graph[T] + start T + queue []T + visited map[T]bool + head int +} + +func NewBreadthFirstIterator[T comparable](g graph.Graph[T], start T) Iterator[T] { + return newBreadthFirstIterator[T](g, start) +} + +func newBreadthFirstIterator[T comparable](g graph.Graph[T], start T) *breadthFirstIterator[T] { + return &breadthFirstIterator[T]{ + graph: g, + start: start, + queue: []T{start}, + visited: make(map[T]bool), + head: -1, + } +} + +func (d *breadthFirstIterator[T]) HasNext() bool { + return d.head < len(d.queue)-1 +} + +func (d *breadthFirstIterator[T]) Next() *graph.Vertex[T] { + d.head++ + + // get the next vertex from the queue + currentNode := d.graph.GetVertexByID(d.queue[d.head]) + + // mark the vertex as visited + + // add unvisited neighbors to the queue + neighbors := currentNode.Neighbors() + for _, neighbor := range neighbors { + if !d.visited[neighbor.Label()] { + d.visited[neighbor.Label()] = true + d.queue = append(d.queue, neighbor.Label()) + } + } + + return currentNode +} + +func (d *breadthFirstIterator[T]) Iterate(f func(v *graph.Vertex[T]) error) error { + for d.HasNext() { + if err := f(d.Next()); err != nil { + return err + } + } + + return nil +} + +func (d *breadthFirstIterator[T]) Reset() { + d.queue = []T{d.start} + d.head = -1 + d.visited = make(map[T]bool) +} diff --git a/graph/traverse/bf_iterator_test.go b/graph/traverse/bf_iterator_test.go new file mode 100644 index 0000000..9f6bf65 --- /dev/null +++ b/graph/traverse/bf_iterator_test.go @@ -0,0 +1,83 @@ +package traverse + +import ( + "reflect" + "testing" + + "github.com/hmdsefi/gograph/graph" +) + +func TestBreadthFirstIterator(t *testing.T) { + // Create a new graph + g := graph.NewDirectedGraph[string]() + + // the example graph + // A -> B -> C + // | | | + // v v v + // D -> E -> F + + vertices := map[string]*graph.Vertex[string]{ + "A": g.AddVertexByLabel("A"), + "B": g.AddVertexByLabel("B"), + "C": g.AddVertexByLabel("C"), + "D": g.AddVertexByLabel("D"), + "E": g.AddVertexByLabel("E"), + "F": g.AddVertexByLabel("F"), + } + + // Add some edges + _, _ = g.AddEdge(vertices["A"], vertices["B"]) + _, _ = g.AddEdge(vertices["A"], vertices["D"]) + _, _ = g.AddEdge(vertices["B"], vertices["C"]) + _, _ = g.AddEdge(vertices["B"], vertices["E"]) + _, _ = g.AddEdge(vertices["C"], vertices["F"]) + _, _ = g.AddEdge(vertices["D"], vertices["E"]) + _, _ = g.AddEdge(vertices["E"], vertices["F"]) + + // Test depth first iteration + iter := newBreadthFirstIterator[string](g, "A") + expected := []string{"A", "B", "D", "C", "E", "F"} + + for i, label := range expected { + if !iter.HasNext() { + t.Errorf("Expected iter.HasNext() to be true, but it was false for label %s", label) + } + + v := iter.Next() + if v.Label() != expected[i] { + t.Errorf("Expected iter.Next().Label() to be %s, but got %s", expected[i], v.Label()) + } + } + + if iter.HasNext() { + t.Error("Expected iter.HasNext() to be false, but it was true") + } + + // test the Reset method + iter.Reset() + if !iter.HasNext() { + t.Error("Expected iter.HasNext() to be true, but it was false after reset") + } + + v := iter.Next() + if v.Label() != "A" { + t.Errorf("Expected iter.Next().Label() to be %s, but got %s", "A", v.Label()) + } + + // test Iterate method + iter.Reset() + var ordered []string + err := iter.Iterate(func(vertex *graph.Vertex[string]) error { + ordered = append(ordered, vertex.Label()) + return nil + }) + if err != nil { + t.Errorf("Expect iter.Iterate(func) returns no error, but got one %s", err) + } + + if !reflect.DeepEqual(expected, ordered) { + t.Errorf("Expect same vertex order, but got different one expected: %v, actual: %v", + expected, ordered) + } +} diff --git a/graph/traverse/dfs_iterator.go b/graph/traverse/df_iterator.go similarity index 92% rename from graph/traverse/dfs_iterator.go rename to graph/traverse/df_iterator.go index dc73fb0..3887086 100644 --- a/graph/traverse/dfs_iterator.go +++ b/graph/traverse/df_iterator.go @@ -23,11 +23,11 @@ func newDepthFirstIterator[T comparable](g graph.Graph[T], start T) *depthFirstI } func (d *depthFirstIterator[T]) HasNext() bool { - return len(d.stack) != 0 + return len(d.stack) > 0 } func (d *depthFirstIterator[T]) Next() *graph.Vertex[T] { - // get the next vertex from the stack + // get the next vertex from the queue label := d.stack[len(d.stack)-1] d.stack = d.stack[:len(d.stack)-1] currentNode := d.graph.GetVertexByID(label) @@ -35,7 +35,7 @@ func (d *depthFirstIterator[T]) Next() *graph.Vertex[T] { // mark the vertex as visited d.visited[label] = true - // add unvisited neighbors to the stack + // add unvisited neighbors to the queue neighbors := currentNode.Neighbors() for _, neighbor := range neighbors { if !d.visited[neighbor.Label()] { diff --git a/graph/traverse/dfs_iterator_test.go b/graph/traverse/df_iterator_test.go similarity index 100% rename from graph/traverse/dfs_iterator_test.go rename to graph/traverse/df_iterator_test.go diff --git a/graph/traverse/iterate.go b/graph/traverse/iterator.go similarity index 100% rename from graph/traverse/iterate.go rename to graph/traverse/iterator.go