diff --git a/graphs_trees/graph_dfs/dfs_challenge.ipynb b/graphs_trees/graph_dfs/dfs_challenge.ipynb
index 6732d68..6fda702 100644
--- a/graphs_trees/graph_dfs/dfs_challenge.ipynb
+++ b/graphs_trees/graph_dfs/dfs_challenge.ipynb
@@ -104,9 +104,11 @@
    },
    "outputs": [],
    "source": [
-    "def dfs(root, visit_func):\n",
-    "    # TODO: Implement me\n",
-    "    pass"
+    "class GraphDfs(Graph):\n",
+    "\n",
+    "    def dfs(self, root, visit_func):\n",
+    "        # TODO: Implement me\n",
+    "        pass"
    ]
   },
   {
@@ -153,7 +155,7 @@
     "\n",
     "    def test_dfs(self):\n",
     "        nodes = []\n",
-    "        graph = Graph()\n",
+    "        graph = GraphDfs()\n",
     "        for id in range(0, 6):\n",
     "            nodes.append(graph.add_node(id))\n",
     "        graph.add_edge(0, 1, 5)\n",
@@ -164,7 +166,7 @@
     "        graph.add_edge(2, 1, 6)\n",
     "        graph.add_edge(3, 2, 7)\n",
     "        graph.add_edge(3, 4, 8)\n",
-    "        dfs(nodes[0], self.results.add_result)\n",
+    "        graph.dfs(nodes[0], self.results.add_result)\n",
     "        assert_equal(str(self.results), \"[0, 1, 3, 2, 4, 5]\")\n",
     "\n",
     "        print('Success: test_dfs')\n",
diff --git a/graphs_trees/graph_dfs/dfs_solution.ipynb b/graphs_trees/graph_dfs/dfs_solution.ipynb
index f0ca00d..7fbb3e3 100644
--- a/graphs_trees/graph_dfs/dfs_solution.ipynb
+++ b/graphs_trees/graph_dfs/dfs_solution.ipynb
@@ -110,14 +110,16 @@
    },
    "outputs": [],
    "source": [
-    "def dfs(root, visit_func):\n",
-    "    if root is None:\n",
-    "        return\n",
-    "    visit_func(root)\n",
-    "    root.visit_state = State.visited\n",
-    "    for node in root.adjacent:\n",
-    "        if node.visit_state == State.unvisited:\n",
-    "            dfs(node, visit_func)"
+    "class GraphDfs(Graph):\n",
+    "\n",
+    "    def dfs(self, root, visit_func):\n",
+    "        if root is None:\n",
+    "            return\n",
+    "        visit_func(root)\n",
+    "        root.visit_state = State.visited\n",
+    "        for node in root.adj_nodes.values():\n",
+    "            if node.visit_state == State.unvisited:\n",
+    "                self.dfs(node, visit_func)"
    ]
   },
   {
@@ -165,7 +167,7 @@
     "\n",
     "    def test_dfs(self):\n",
     "        nodes = []\n",
-    "        graph = Graph()\n",
+    "        graph = GraphDfs()\n",
     "        for id in range(0, 6):\n",
     "            nodes.append(graph.add_node(id))\n",
     "        graph.add_edge(0, 1, 5)\n",
@@ -176,7 +178,7 @@
     "        graph.add_edge(2, 1, 6)\n",
     "        graph.add_edge(3, 2, 7)\n",
     "        graph.add_edge(3, 4, 8)\n",
-    "        dfs(nodes[0], self.results.add_result)\n",
+    "        graph.dfs(nodes[0], self.results.add_result)\n",
     "        assert_equal(str(self.results), \"[0, 1, 3, 2, 4, 5]\")\n",
     "\n",
     "        print('Success: test_dfs')\n",
diff --git a/graphs_trees/graph_dfs/test_dfs.py b/graphs_trees/graph_dfs/test_dfs.py
index 2830387..8e1c50e 100644
--- a/graphs_trees/graph_dfs/test_dfs.py
+++ b/graphs_trees/graph_dfs/test_dfs.py
@@ -8,7 +8,7 @@ class TestDfs(object):
 
     def test_dfs(self):
         nodes = []
-        graph = Graph()
+        graph = GraphDfs()
         for id in range(0, 6):
             nodes.append(graph.add_node(id))
         graph.add_edge(0, 1, 5)
@@ -19,7 +19,7 @@ class TestDfs(object):
         graph.add_edge(2, 1, 6)
         graph.add_edge(3, 2, 7)
         graph.add_edge(3, 4, 8)
-        dfs(nodes[0], self.results.add_result)
+        graph.dfs(nodes[0], self.results.add_result)
         assert_equal(str(self.results), "[0, 1, 3, 2, 4, 5]")
 
         print('Success: test_dfs')