diff --git a/stacks_queues/sort_stack/sort_stack_challenge.ipynb b/stacks_queues/sort_stack/sort_stack_challenge.ipynb
index 1132e50..3dd087b 100644
--- a/stacks_queues/sort_stack/sort_stack_challenge.ipynb
+++ b/stacks_queues/sort_stack/sort_stack_challenge.ipynb
@@ -123,26 +123,25 @@
     "\n",
     "class TestSortStack(object):\n",
     "\n",
-    "    def get_sorted_stack(self, numbers):\n",
-    "        stack = MyStack()\n",
+    "    def get_sorted_stack(self, stack, numbers):\n",
     "        for x in numbers:\n",
     "            stack.push(x)\n",
     "        sorted_stack = stack.sort()\n",
     "        return sorted_stack\n",
     "\n",
-    "    def test_sort_stack(self):\n",
+    "    def test_sort_stack(self, stack):\n",
     "        print('Test: Empty stack')\n",
-    "        sorted_stack = self.get_sorted_stack([])\n",
+    "        sorted_stack = self.get_sorted_stack(stack, [])\n",
     "        assert_equal(sorted_stack.pop(), None)\n",
     "\n",
     "        print('Test: One element stack')\n",
-    "        sorted_stack = self.get_sorted_stack([1])\n",
+    "        sorted_stack = self.get_sorted_stack(stack, [1])\n",
     "        assert_equal(sorted_stack.pop(), 1)\n",
     "\n",
     "        print('Test: Two or more element stack (general case)')\n",
     "        num_items = 10\n",
     "        numbers = [randint(0, 10) for x in range(num_items)]\n",
-    "        sorted_stack = self.get_sorted_stack(numbers)\n",
+    "        sorted_stack = self.get_sorted_stack(stack, numbers)\n",
     "        sorted_numbers = []\n",
     "        for _ in range(num_items):\n",
     "            sorted_numbers.append(sorted_stack.pop())\n",
@@ -153,7 +152,8 @@
     "\n",
     "def main():\n",
     "    test = TestSortStack()\n",
-    "    test.test_sort_stack()\n",
+    "    test.test_sort_stack(MyStack())\n",
+    "    test.test_sort_stack(MyStackSimplified())\n",
     "\n",
     "\n",
     "if __name__ == '__main__':\n",
@@ -172,21 +172,21 @@
  ],
  "metadata": {
   "kernelspec": {
-   "display_name": "Python 2",
+   "display_name": "Python 3",
    "language": "python",
-   "name": "python2"
+   "name": "python3"
   },
   "language_info": {
    "codemirror_mode": {
     "name": "ipython",
-    "version": 2
+    "version": 3
    },
    "file_extension": ".py",
    "mimetype": "text/x-python",
    "name": "python",
    "nbconvert_exporter": "python",
-   "pygments_lexer": "ipython2",
-   "version": "2.7.10"
+   "pygments_lexer": "ipython3",
+   "version": "3.4.3"
   }
  },
  "nbformat": 4,
diff --git a/stacks_queues/sort_stack/sort_stack_solution.ipynb b/stacks_queues/sort_stack/sort_stack_solution.ipynb
index 4e51e64..7ac8a0d 100644
--- a/stacks_queues/sort_stack/sort_stack_solution.ipynb
+++ b/stacks_queues/sort_stack/sort_stack_solution.ipynb
@@ -105,6 +105,36 @@
     "        buff = MyStack()\n",
     "        while not self.is_empty():\n",
     "            temp = self.pop()\n",
+    "            if buff.is_empty() or temp >= buff.peek():\n",
+    "                buff.push(temp)\n",
+    "            else:\n",
+    "                while not buff.is_empty() and buff.peek() > temp:\n",
+    "                    self.push(buff.pop())\n",
+    "                buff.push(temp)\n",
+    "        return buff"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "The solution can be further simplified:"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 3,
+   "metadata": {
+    "collapsed": true
+   },
+   "outputs": [],
+   "source": [
+    "class MyStackSimplified(Stack):\n",
+    "\n",
+    "    def sort(self):\n",
+    "        buff = MyStack()\n",
+    "        while not self.is_empty():\n",
+    "            temp = self.pop()\n",
     "            while not buff.is_empty() and buff.peek() > temp:\n",
     "                self.push(buff.pop())\n",
     "            buff.push(temp)\n",
@@ -121,7 +151,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 3,
+   "execution_count": 4,
    "metadata": {
     "collapsed": false
    },
@@ -142,26 +172,25 @@
     "\n",
     "class TestSortStack(object):\n",
     "\n",
-    "    def get_sorted_stack(self, numbers):\n",
-    "        stack = MyStack()\n",
+    "    def get_sorted_stack(self, stack, numbers):\n",
     "        for x in numbers:\n",
     "            stack.push(x)\n",
     "        sorted_stack = stack.sort()\n",
     "        return sorted_stack\n",
     "\n",
-    "    def test_sort_stack(self):\n",
+    "    def test_sort_stack(self, stack):\n",
     "        print('Test: Empty stack')\n",
-    "        sorted_stack = self.get_sorted_stack([])\n",
+    "        sorted_stack = self.get_sorted_stack(stack, [])\n",
     "        assert_equal(sorted_stack.pop(), None)\n",
     "\n",
     "        print('Test: One element stack')\n",
-    "        sorted_stack = self.get_sorted_stack([1])\n",
+    "        sorted_stack = self.get_sorted_stack(stack, [1])\n",
     "        assert_equal(sorted_stack.pop(), 1)\n",
     "\n",
     "        print('Test: Two or more element stack (general case)')\n",
     "        num_items = 10\n",
     "        numbers = [randint(0, 10) for x in range(num_items)]\n",
-    "        sorted_stack = self.get_sorted_stack(numbers)\n",
+    "        sorted_stack = self.get_sorted_stack(stack, numbers)\n",
     "        sorted_numbers = []\n",
     "        for _ in range(num_items):\n",
     "            sorted_numbers.append(sorted_stack.pop())\n",
@@ -172,7 +201,8 @@
     "\n",
     "def main():\n",
     "    test = TestSortStack()\n",
-    "    test.test_sort_stack()\n",
+    "    test.test_sort_stack(MyStack())\n",
+    "    test.test_sort_stack(MyStackSimplified())\n",
     "\n",
     "\n",
     "if __name__ == '__main__':\n",
@@ -181,7 +211,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 4,
+   "execution_count": 5,
    "metadata": {
     "collapsed": false
    },
@@ -190,6 +220,10 @@
      "name": "stdout",
      "output_type": "stream",
      "text": [
+      "Test: Empty stack\n",
+      "Test: One element stack\n",
+      "Test: Two or more element stack (general case)\n",
+      "Success: test_sort_stack\n",
       "Test: Empty stack\n",
       "Test: One element stack\n",
       "Test: Two or more element stack (general case)\n",
@@ -204,21 +238,21 @@
  ],
  "metadata": {
   "kernelspec": {
-   "display_name": "Python 2",
+   "display_name": "Python 3",
    "language": "python",
-   "name": "python2"
+   "name": "python3"
   },
   "language_info": {
    "codemirror_mode": {
     "name": "ipython",
-    "version": 2
+    "version": 3
    },
    "file_extension": ".py",
    "mimetype": "text/x-python",
    "name": "python",
    "nbconvert_exporter": "python",
-   "pygments_lexer": "ipython2",
-   "version": "2.7.10"
+   "pygments_lexer": "ipython3",
+   "version": "3.4.3"
   }
  },
  "nbformat": 4,
diff --git a/stacks_queues/sort_stack/test_sort_stack.py b/stacks_queues/sort_stack/test_sort_stack.py
index e117053..f86e931 100644
--- a/stacks_queues/sort_stack/test_sort_stack.py
+++ b/stacks_queues/sort_stack/test_sort_stack.py
@@ -4,26 +4,25 @@ from nose.tools import assert_equal
 
 class TestSortStack(object):
 
-    def get_sorted_stack(self, numbers):
-        stack = MyStack()
+    def get_sorted_stack(self, stack, numbers):
         for x in numbers:
             stack.push(x)
         sorted_stack = stack.sort()
         return sorted_stack
 
-    def test_sort_stack(self):
+    def test_sort_stack(self, stack):
         print('Test: Empty stack')
-        sorted_stack = self.get_sorted_stack([])
+        sorted_stack = self.get_sorted_stack(stack, [])
         assert_equal(sorted_stack.pop(), None)
 
         print('Test: One element stack')
-        sorted_stack = self.get_sorted_stack([1])
+        sorted_stack = self.get_sorted_stack(stack, [1])
         assert_equal(sorted_stack.pop(), 1)
 
         print('Test: Two or more element stack (general case)')
         num_items = 10
         numbers = [randint(0, 10) for x in range(num_items)]
-        sorted_stack = self.get_sorted_stack(numbers)
+        sorted_stack = self.get_sorted_stack(stack, numbers)
         sorted_numbers = []
         for _ in range(num_items):
             sorted_numbers.append(sorted_stack.pop())
@@ -34,7 +33,8 @@ class TestSortStack(object):
 
 def main():
     test = TestSortStack()
-    test.test_sort_stack()
+    test.test_sort_stack(MyStack())
+    test.test_sort_stack(MyStackSimplified())
 
 
 if __name__ == '__main__':