fix bug when removing vertex from directed graph

earlygrey · Sep 25, 2022 · 0865020 · 0865020
1 parent 23a1488
commit 0865020
Show file tree

Hide file tree

Showing 2 changed files with 59 additions and 7 deletions.
diff --git a/src/main/java/space/earlygrey/simplegraphs/Graph.java b/src/main/java/space/earlygrey/simplegraphs/Graph.java
@@ -133,20 +133,26 @@ public void addVertices(V... vertices) {
     public boolean removeVertex(V v) {
         Node<V> existing = nodeMap.remove(v);
         if (existing == null) return false;
-        for (int i = existing.getOutEdges().size() - 1; i >= 0; i--) {
-            removeConnection(existing.getOutEdges().get(i).b, existing);
-        }
-        existing.disconnect();
+        disconnect(existing);
         return true;
     }
 
     public void disconnect(V v) {
         Node<V> existing = nodeMap.get(v);
         if (existing == null) Errors.throwVertexNotInGraphVertexException(false);
-        for (int i = existing.getOutEdges().size() - 1; i >= 0; i--) {
-            removeConnection(existing.getOutEdges().get(i).b, existing);
+        disconnect(existing);
+    }
+
+    protected void disconnect(Node<V> node) {
+        for (int i = node.getOutEdges().size() - 1; i >= 0; i--) {
+            removeConnection(node, node.getOutEdges().get(i).b);
+        }
+        if (node.getInEdges() != null) {
+            for (int i = node.getInEdges().size() - 1; i >= 0; i--) {
+                removeConnection(node.getInEdges().get(i).a, node);
+            }
         }
-        existing.disconnect();
+        node.disconnect();
     }
 
     /**

diff --git a/src/test/java/space/earlygrey/simplegraphs/GraphTest.java b/src/test/java/space/earlygrey/simplegraphs/GraphTest.java
@@ -127,6 +127,52 @@ public void verticesCanBeAddedAndRemoved() {
         }
     }
 
+    @Test
+    public void vertexanBeRemovedFromDirectedGraph() {
+        DirectedGraph<Integer> graph = new DirectedGraph<>();
+
+        int n = 10;
+
+        for (int i = 0; i < n; i++) {
+            graph.addVertex(i);
+        }
+
+        for (int i = 0; i < n-1; i++) {
+            graph.addEdge(i, i+1);
+        }
+
+        assertEquals(n, graph.size());
+        assertEquals(n-1, graph.getEdgeCount());
+
+        graph.removeVertex(n/2);
+
+        assertEquals(n-1, graph.size());
+        assertEquals(n-3, graph.getEdgeCount());
+    }
+
+    @Test
+    public void vertexCanBeDisconnectedFromDirectedGraph() {
+        DirectedGraph<Integer> graph = new DirectedGraph<>();
+
+        int n = 10;
+
+        for (int i = 0; i < n; i++) {
+            graph.addVertex(i);
+        }
+
+        for (int i = 0; i < n-1; i++) {
+            graph.addEdge(i, i+1);
+        }
+
+        assertEquals(n, graph.size());
+        assertEquals(n-1, graph.getEdgeCount());
+
+        graph.disconnect(n/2);
+
+        assertEquals(n, graph.size());
+        assertEquals(n-3, graph.getEdgeCount());
+    }
+
     @Test
     public void edgesCanBeAddedAndRemoved() {
         int n = 5;