Merge pull request #3767 from juliantaylor/mom5-improve

Mom5 selection improvements

2 files changed, 13 insertions, 6 deletions

diff --git a/numpy/core/src/npysort/selection.c.src b/numpy/core/src/npysort/selection.c.src
index b11753367..073b5847f 100644
--- 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
index 77da5543d..ee773ef4a 100644
--- a/numpy/core/tests/test_multiarray.py
+++ b/numpy/core/tests/test_multiarray.py
@@ -1141,6 +1141,17 @@ class TestMethods(TestCase):
             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)