Merge pull request numpy#3767 from juliantaylor/mom5-improve

Mom5 selection improvements
jcrabtree · Sep 22, 2013 · 0b96b2d · 0b96b2d
2 parents fffeee6 + eb8991e
commit 0b96b2d
Show file tree

Hide file tree

Showing 2 changed files with 13 additions and 6 deletions.
diff --git a/numpy/core/src/npysort/selection.c.src b/numpy/core/src/npysort/selection.c.src
@@ -176,12 +176,8 @@ static npy_intp @name@median5_@suff@(
         }
     }
     else {
-        if (@TYPE@_LT(v[IDX(2)], v[IDX(1)])) {
-            return 1;
-        }
-        else {
-            return 2;
-        }
+        /* v[1] and v[2] swapped into order above */
+        return 2;
     }
 }
 

diff --git a/numpy/core/tests/test_multiarray.py b/numpy/core/tests/test_multiarray.py
@@ -1141,6 +1141,17 @@ def test_partition(self):
             assert_array_equal(d[np.argpartition(d, -6, kind=k)],
                                np.partition(d, 41, kind=k))
 
+            # median of 3 killer, O(n^2) on pure median 3 pivot quickselect
+            # exercises the median of median of 5 code used to keep O(n)
+            d = np.arange(1000000)
+            x = np.roll(d, d.size // 2)
+            mid = x.size // 2 + 1
+            assert_equal(np.partition(x, mid)[mid], mid)
+            d = np.arange(1000001)
+            x = np.roll(d, d.size // 2 + 1)
+            mid = x.size // 2 + 1
+            assert_equal(np.partition(x, mid)[mid], mid)
+
             # equal elements
             d = np.arange((47)) % 7
             tgt = np.sort(np.arange((47)) % 7)