diff --git a/graphs_trees/graph/graph.py b/graphs_trees/graph/graph.py
index d7e3e96..03f0a62 100644
--- a/graphs_trees/graph/graph.py
+++ b/graphs_trees/graph/graph.py
@@ -30,14 +30,11 @@ class Node:
         self.adj_nodes[neighbor.key] = neighbor
 
     def remove_neighbor(self, neighbor):
-        if neighbor is None:
-            raise Exception('Invalid neighbor')
-        if neighbor.key in self.adj_nodes:
-            neighbor.incoming_edges -= 1
-            del self.adj_weights[neighbor.key]
-            del self.adj_nodes[neighbor.key]
-        else:
+        if neighbor is None or neighbor.key not in self.adj_nodes:
             raise Exception('Invalid neighbor')
+        neighbor.incoming_edges -= 1
+        del self.adj_weights[neighbor.key]
+        del self.adj_nodes[neighbor.key]
 
 
 class Graph:
diff --git a/graphs_trees/graph/graph_solution.ipynb b/graphs_trees/graph/graph_solution.ipynb
index e5b982b..f30ddf4 100644
--- a/graphs_trees/graph/graph_solution.ipynb
+++ b/graphs_trees/graph/graph_solution.ipynb
@@ -194,14 +194,11 @@
     "        self.adj_nodes[neighbor.key] = neighbor\n",
     "\n",
     "    def remove_neighbor(self, neighbor):\n",
-    "        if neighbor is None:\n",
-    "            raise Exception('Invalid neighbor')\n",
-    "        if neighbor.key in self.adj_nodes:\n",
-    "            neighbor.incoming_edges -= 1\n",
-    "            del self.adj_weights[neighbor.key]\n",
-    "            del self.adj_nodes[neighbor.key]\n",
-    "        else:\n",
+    "        if neighbor is None or neighbor.key not in self.adj_nodes:\n",
     "            raise Exception('Invalid neighbor')\n",
+    "        neighbor.incoming_edges -= 1\n",
+    "        del self.adj_weights[neighbor.key]\n",
+    "        del self.adj_nodes[neighbor.key]\n",
     "\n",
     "\n",
     "class Graph:\n",