diff --git a/graphs_trees/bst_successor/bst_successor_solution.ipynb b/graphs_trees/bst_successor/bst_successor_solution.ipynb
index b756c93..4ee62ff 100644
--- a/graphs_trees/bst_successor/bst_successor_solution.ipynb
+++ b/graphs_trees/bst_successor/bst_successor_solution.ipynb
@@ -115,21 +115,24 @@
     "\n",
     "    def get_next(self, node):\n",
     "        if node is None:\n",
-    "            raise Exception('Invalid input node')\n",
+    "            raise TypeError('node cannot be None')\n",
     "        if node.right is not None:\n",
     "            return self._left_most(node.right)\n",
-    "        return self._next_ancestor(node)\n",
+    "        else:\n",
+    "            return self._next_ancestor(node)\n",
     "\n",
     "    def _left_most(self, node):\n",
     "        if node.left is not None:\n",
     "            return self._left_most(node.left)\n",
-    "        return node.data\n",
+    "        else:\n",
+    "            return node.data\n",
     "\n",
     "    def _next_ancestor(self, node):\n",
     "        if node.parent is not None:\n",
     "            if node.parent.data > node.data:\n",
     "                return node.parent.data\n",
-    "            return self._next_ancestor(node.parent)\n",
+    "            else:\n",
+    "                return self._next_ancestor(node.parent)\n",
     "        # We reached the root, the original input node\n",
     "        # must be the largest element in the tree.\n",
     "        return None"