diff --git a/dag.go b/dag.go index 6682772..c7c027a 100644 --- a/dag.go +++ b/dag.go @@ -118,6 +118,11 @@ func (d *DAG) DeleteEdge(tailVertex *Vertex, headVertex *Vertex) error { tailVertex.Children.Remove(childVertex) } } + for _, parentVertex := range headVertex.Parents.Values() { + if parentVertex == tailVertex { + headVertex.Parents.Remove(parentVertex) + } + } return nil } diff --git a/dag_test.go b/dag_test.go index f5233d6..9f8ebac 100644 --- a/dag_test.go +++ b/dag_test.go @@ -183,6 +183,16 @@ func TestDAG_DeleteEdge(t *testing.T) { if size != 0 { t.Fatalf("Dag expected to have 0 edges but got %d", size) } + + outDegree := vertex1.OutDegree() + if outDegree != 0 { + t.Fatalf("Vertex %s expected to have 0 outgoing edges but got %d", vertex1.ID, outDegree) + } + inDegree := vertex2.InDegree() + if inDegree != 0 { + t.Fatalf("Vertex %s expected to have 0 incoming edges but got %d", vertex2.ID, inDegree) + } + } func TestDAG_GetVertex(t *testing.T) {