diff --git a/graphs_trees/tree_height/height.py b/graphs_trees/tree_height/height.py
index 581900b..d95f80c 100644
--- a/graphs_trees/tree_height/height.py
+++ b/graphs_trees/tree_height/height.py
@@ -1,5 +1,7 @@
-def height(node):
-    if node is None:
-        return 0
-    return 1 + max(height(node.left), 
-                   height(node.right))
\ No newline at end of file
+class BstHeight(Bst):
+
+    def height(self, node):
+        if node is None:
+            return 0
+        return 1 + max(self.height(node.left),
+                       self.height(node.right))
\ No newline at end of file
diff --git a/graphs_trees/tree_height/height_challenge.ipynb b/graphs_trees/tree_height/height_challenge.ipynb
index b2041c7..f23465d 100644
--- a/graphs_trees/tree_height/height_challenge.ipynb
+++ b/graphs_trees/tree_height/height_challenge.ipynb
@@ -88,9 +88,11 @@
    },
    "outputs": [],
    "source": [
-    "def height(node):\n",
-    "    # TODO: Implement me\n",
-    "    pass"
+    "class BstHeight(Bst):\n",
+    "\n",
+    "    def height(self, node):\n",
+    "        # TODO: Implement me\n",
+    "        pass"
    ]
   },
   {
@@ -122,13 +124,13 @@
     "class TestHeight(object):\n",
     "\n",
     "    def test_height(self):\n",
-    "        root = Node(5)\n",
-    "        assert_equal(height(root), 1)\n",
-    "        insert(root, 2)\n",
-    "        insert(root, 8)\n",
-    "        insert(root, 1)\n",
-    "        insert(root, 3)\n",
-    "        assert_equal(height(root), 3)\n",
+    "        bst = BstHeight(Node(5))\n",
+    "        assert_equal(bst.height(bst.root), 1)\n",
+    "        bst.insert(2)\n",
+    "        bst.insert(8)\n",
+    "        bst.insert(1)\n",
+    "        bst.insert(3)\n",
+    "        assert_equal(bst.height(bst.root), 3)\n",
     "\n",
     "        print('Success: test_height')\n",
     "\n",
diff --git a/graphs_trees/tree_height/height_solution.ipynb b/graphs_trees/tree_height/height_solution.ipynb
index 7d8a34b..3a70b93 100644
--- a/graphs_trees/tree_height/height_solution.ipynb
+++ b/graphs_trees/tree_height/height_solution.ipynb
@@ -91,33 +91,15 @@
    "metadata": {
     "collapsed": false
    },
-   "outputs": [
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "Overwriting height.py\n"
-     ]
-    }
-   ],
-   "source": [
-    "%%writefile height.py\n",
-    "def height(node):\n",
-    "    if node is None:\n",
-    "        return 0\n",
-    "    return 1 + max(height(node.left), \n",
-    "                   height(node.right))"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 3,
-   "metadata": {
-    "collapsed": true
-   },
    "outputs": [],
    "source": [
-    "%run height.py"
+    "class BstHeight(Bst):\n",
+    "\n",
+    "    def height(self, node):\n",
+    "        if node is None:\n",
+    "            return 0\n",
+    "        return 1 + max(self.height(node.left),\n",
+    "                       self.height(node.right))"
    ]
   },
   {
@@ -129,7 +111,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 4,
+   "execution_count": 3,
    "metadata": {
     "collapsed": false
    },
@@ -150,13 +132,13 @@
     "class TestHeight(object):\n",
     "\n",
     "    def test_height(self):\n",
-    "        root = Node(5)\n",
-    "        assert_equal(height(root), 1)\n",
-    "        insert(root, 2)\n",
-    "        insert(root, 8)\n",
-    "        insert(root, 1)\n",
-    "        insert(root, 3)\n",
-    "        assert_equal(height(root), 3)\n",
+    "        bst = BstHeight(Node(5))\n",
+    "        assert_equal(bst.height(bst.root), 1)\n",
+    "        bst.insert(2)\n",
+    "        bst.insert(8)\n",
+    "        bst.insert(1)\n",
+    "        bst.insert(3)\n",
+    "        assert_equal(bst.height(bst.root), 3)\n",
     "\n",
     "        print('Success: test_height')\n",
     "\n",
@@ -172,7 +154,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 5,
+   "execution_count": 4,
    "metadata": {
     "collapsed": false
    },
diff --git a/graphs_trees/tree_height/test_height.py b/graphs_trees/tree_height/test_height.py
index b1f2f68..fbb45cb 100644
--- a/graphs_trees/tree_height/test_height.py
+++ b/graphs_trees/tree_height/test_height.py
@@ -4,13 +4,13 @@ from nose.tools import assert_equal
 class TestHeight(object):
 
     def test_height(self):
-        root = Node(5)
-        assert_equal(height(root), 1)
-        insert(root, 2)
-        insert(root, 8)
-        insert(root, 1)
-        insert(root, 3)
-        assert_equal(height(root), 3)
+        bst = BstHeight(Node(5))
+        assert_equal(bst.height(bst.root), 1)
+        bst.insert(2)
+        bst.insert(8)
+        bst.insert(1)
+        bst.insert(3)
+        assert_equal(bst.height(bst.root), 3)
 
         print('Success: test_height')