diff --git a/arrays_strings/compress/compress_solution.ipynb b/arrays_strings/compress/compress_solution.ipynb
index 3f14c6e..43e5bc6 100644
--- a/arrays_strings/compress/compress_solution.ipynb
+++ b/arrays_strings/compress/compress_solution.ipynb
@@ -108,7 +108,7 @@
     "            result += prev_char + (str(count) if count > 1 else '')\n",
     "            prev_char = char\n",
     "            count = 1\n",
-    "    result += prev_char + str(count)\n",
+    "    result += prev_char + (str(count) if count > 1 else '')\n",
     "    return result if len(result) < len(string) else string"
    ]
   },
@@ -145,7 +145,7 @@
     "        assert_equal(func(None), None)\n",
     "        assert_equal(func(''), '')\n",
     "        assert_equal(func('AABBCC'), 'AABBCC')\n",
-    "        assert_equal(func('AAABCCDDDD'), 'A3BC2D4')\n",
+    "        assert_equal(func('AAABCCDDDDE'), 'A3BC2D4E')\n",
     "        print('Success: test_compress')\n",
     "\n",
     "\n",