diff --git a/sorting_searching/merge_sort/merge_sort_challenge.ipynb b/sorting_searching/merge_sort/merge_sort_challenge.ipynb
index 0b7506b..33de82e 100644
--- a/sorting_searching/merge_sort/merge_sort_challenge.ipynb
+++ b/sorting_searching/merge_sort/merge_sort_challenge.ipynb
@@ -37,6 +37,10 @@
     "* Is a naiive solution sufficient?\n",
     "    * Yes\n",
     "* Are duplicates allowed?\n",
+    "    * Yes\n",
+    "* Can we assume the input is valid?\n",
+    "    * No\n",
+    "* Can we assume this fits memory?\n",
     "    * Yes"
    ]
   },
@@ -46,7 +50,7 @@
    "source": [
     "## Test Cases\n",
     "\n",
-    "* None -> None\n",
+    "* None -> Exception\n",
     "* Empty input -> []\n",
     "* One element -> [element]\n",
     "* Two or more elements\n",
@@ -77,9 +81,11 @@
    },
    "outputs": [],
    "source": [
-    "def merge_sort(data):\n",
-    "    # TODO: Implement me\n",
-    "    pass"
+    "class MergeSort(object):\n",
+    "\n",
+    "    def sort(self, data):\n",
+    "        # TODO: Implement me\n",
+    "        pass"
    ]
   },
   {
@@ -102,30 +108,26 @@
    "outputs": [],
    "source": [
     "# %load test_merge_sort.py\n",
-    "from nose.tools import assert_equal\n",
+    "from nose.tools import assert_equal, assert_raises\n",
     "\n",
     "\n",
     "class TestMergeSort(object):\n",
+    "\n",
     "    def test_merge_sort(self):\n",
+    "        merge_sort = MergeSort()\n",
+    "\n",
     "        print('None input')\n",
-    "        data = None\n",
-    "        sorted_data = merge_sort(data)\n",
-    "        assert_equal(sorted_data, None)\n",
+    "        assert_raises(TypeError, merge_sort.sort, None)\n",
     "\n",
     "        print('Empty input')\n",
-    "        data = []\n",
-    "        sorted_data = merge_sort(data)\n",
-    "        assert_equal(sorted_data, [])\n",
+    "        assert_equal(merge_sort.sort([]), [])\n",
     "\n",
     "        print('One element')\n",
-    "        data = [5]\n",
-    "        sorted_data = merge_sort(data)\n",
-    "        assert_equal(sorted_data, [5])\n",
+    "        assert_equal(merge_sort.sort([5]), [5])\n",
     "\n",
     "        print('Two or more elements')\n",
     "        data = [5, 1, 7, 2, 6, -3, 5, 7, -1]\n",
-    "        sorted_data = merge_sort(data)\n",
-    "        assert_equal(sorted_data, sorted(data))\n",
+    "        assert_equal(merge_sort.sort(data), sorted(data))\n",
     "\n",
     "        print('Success: test_merge_sort')\n",
     "\n",
diff --git a/sorting_searching/merge_sort/merge_sort_solution.ipynb b/sorting_searching/merge_sort/merge_sort_solution.ipynb
index 3a27cdc..1fde348 100644
--- a/sorting_searching/merge_sort/merge_sort_solution.ipynb
+++ b/sorting_searching/merge_sort/merge_sort_solution.ipynb
@@ -36,6 +36,10 @@
     "* Is a naiive solution sufficient?\n",
     "    * Yes\n",
     "* Are duplicates allowed?\n",
+    "    * Yes\n",
+    "* Can we assume the input is valid?\n",
+    "    * No\n",
+    "* Can we assume this fits memory?\n",
     "    * Yes"
    ]
   },
@@ -45,7 +49,7 @@
    "source": [
     "## Test Cases\n",
     "\n",
-    "* None -> None\n",
+    "* None -> Exception\n",
     "* Empty input -> []\n",
     "* One element -> [element]\n",
     "* Two or more elements\n",
@@ -94,37 +98,42 @@
     "from __future__ import division\n",
     "\n",
     "\n",
-    "def merge(left, right):\n",
-    "    l = 0\n",
-    "    r = 0\n",
-    "    result = []\n",
-    "    while l < len(left) and r < len(right):\n",
-    "        if left[l] < right[r]:\n",
+    "class MergeSort(object):\n",
+    "\n",
+    "    def sort(self, data):\n",
+    "        if data is None:\n",
+    "            raise TypeError('data cannot be None')\n",
+    "        return self._sort(data)\n",
+    "\n",
+    "    def _sort(self, data):\n",
+    "        if len(data) < 2:\n",
+    "            return data\n",
+    "        mid = len(data) // 2\n",
+    "        left = data[:mid]\n",
+    "        right = data[mid:]\n",
+    "        left = self._sort(left)\n",
+    "        right = self._sort(right)\n",
+    "        return self._merge(left, right)\n",
+    "\n",
+    "    def _merge(self, left, right):\n",
+    "        l = 0\n",
+    "        r = 0\n",
+    "        result = []\n",
+    "        while l < len(left) and r < len(right):\n",
+    "            if left[l] < right[r]:\n",
+    "                result.append(left[l])\n",
+    "                l += 1\n",
+    "            else:\n",
+    "                result.append(right[r])\n",
+    "                r += 1\n",
+    "        # Copy remaining elements\n",
+    "        while l < len(left):\n",
     "            result.append(left[l])\n",
     "            l += 1\n",
-    "        else:\n",
+    "        while r < len(right):\n",
     "            result.append(right[r])\n",
     "            r += 1\n",
-    "            \n",
-    "    # Copy remaining elements\n",
-    "    while l < len(left):\n",
-    "        result.append(left[l])\n",
-    "        l += 1\n",
-    "    while r < len(right):\n",
-    "        result.append(right[r])\n",
-    "        r += 1\n",
-    "    return result\n",
-    "\n",
-    "\n",
-    "def merge_sort(data):\n",
-    "    if data is None or len(data) < 2:\n",
-    "        return data\n",
-    "    mid = len(data) // 2\n",
-    "    left = data[:mid]\n",
-    "    right = data[mid:]\n",
-    "    left = merge_sort(left)\n",
-    "    right = merge_sort(right)\n",
-    "    return merge(left, right)"
+    "        return result"
    ]
   },
   {
@@ -152,30 +161,26 @@
    ],
    "source": [
     "%%writefile test_merge_sort.py\n",
-    "from nose.tools import assert_equal\n",
+    "from nose.tools import assert_equal, assert_raises\n",
     "\n",
     "\n",
     "class TestMergeSort(object):\n",
+    "\n",
     "    def test_merge_sort(self):\n",
+    "        merge_sort = MergeSort()\n",
+    "\n",
     "        print('None input')\n",
-    "        data = None\n",
-    "        sorted_data = merge_sort(data)\n",
-    "        assert_equal(sorted_data, None)\n",
+    "        assert_raises(TypeError, merge_sort.sort, None)\n",
     "\n",
     "        print('Empty input')\n",
-    "        data = []\n",
-    "        sorted_data = merge_sort(data)\n",
-    "        assert_equal(sorted_data, [])\n",
+    "        assert_equal(merge_sort.sort([]), [])\n",
     "\n",
     "        print('One element')\n",
-    "        data = [5]\n",
-    "        sorted_data = merge_sort(data)\n",
-    "        assert_equal(sorted_data, [5])\n",
+    "        assert_equal(merge_sort.sort([5]), [5])\n",
     "\n",
     "        print('Two or more elements')\n",
     "        data = [5, 1, 7, 2, 6, -3, 5, 7, -1]\n",
-    "        sorted_data = merge_sort(data)\n",
-    "        assert_equal(sorted_data, sorted(data))\n",
+    "        assert_equal(merge_sort.sort(data), sorted(data))\n",
     "\n",
     "        print('Success: test_merge_sort')\n",
     "\n",
diff --git a/sorting_searching/merge_sort/test_merge_sort.py b/sorting_searching/merge_sort/test_merge_sort.py
index 109cb19..e1d8fd9 100644
--- a/sorting_searching/merge_sort/test_merge_sort.py
+++ b/sorting_searching/merge_sort/test_merge_sort.py
@@ -1,27 +1,23 @@
-from nose.tools import assert_equal
+from nose.tools import assert_equal, assert_raises
 
 
 class TestMergeSort(object):
+
     def test_merge_sort(self):
+        merge_sort = MergeSort()
+
         print('None input')
-        data = None
-        sorted_data = merge_sort(data)
-        assert_equal(sorted_data, None)
+        assert_raises(TypeError, merge_sort.sort, None)
 
         print('Empty input')
-        data = []
-        sorted_data = merge_sort(data)
-        assert_equal(sorted_data, [])
+        assert_equal(merge_sort.sort([]), [])
 
         print('One element')
-        data = [5]
-        sorted_data = merge_sort(data)
-        assert_equal(sorted_data, [5])
+        assert_equal(merge_sort.sort([5]), [5])
 
         print('Two or more elements')
         data = [5, 1, 7, 2, 6, -3, 5, 7, -1]
-        sorted_data = merge_sort(data)
-        assert_equal(sorted_data, sorted(data))
+        assert_equal(merge_sort.sort(data), sorted(data))
 
         print('Success: test_merge_sort')