BUG: wrong selection for orders falling into equal ranges

juliantaylor · juliantaylor · commit d6c7a16af4de · 2014-07-04T19:00:59.000+02:00
when orders are selected where the kth element falls into an equal range the the last stored pivot was not the kth element, this leads to losing the ordering of smaller orders as following selection steps can start at index 0 again instead of the at the offset of the last selection. Closes numpygh-4836
diff --git a/numpy/core/src/npysort/selection.c.src b/numpy/core/src/npysort/selection.c.src
@@ -390,7 +390,10 @@ int
         /* move pivot into position */
         SWAP(SORTEE(low), SORTEE(hh));
 
-        store_pivot(hh, kth, pivots, npiv);
+        /* kth pivot stored later */
+        if (hh != kth) {
+            store_pivot(hh, kth, pivots, npiv);
+        }
 
         if (hh >= kth)
             high = hh - 1;
@@ -400,10 +403,11 @@ int
 
     /* two elements */
     if (high == low + 1) {
-        if (@TYPE@_LT(v[IDX(high)], v[IDX(low)]))
+        if (@TYPE@_LT(v[IDX(high)], v[IDX(low)])) {
             SWAP(SORTEE(high), SORTEE(low))
-        store_pivot(low, kth, pivots, npiv);
+        }
     }
+    store_pivot(kth, kth, pivots, npiv);
 
     return 0;
 }
diff --git a/numpy/core/tests/test_multiarray.py b/numpy/core/tests/test_multiarray.py
@@ -1356,6 +1356,12 @@ def test_partition(self):
                 d[i:].partition(0, kind=k)
             assert_array_equal(d, tgt)
 
+            d = np.array([0, 1, 2, 3, 4, 5, 7, 7, 7, 7, 7, 7, 7, 7, 7, 7, 7,
+                          7, 7, 7, 7, 7, 9])
+            kth = [0, 3, 19, 20]
+            assert_equal(np.partition(d, kth, kind=k)[kth], (0, 3, 7, 7))
+            assert_equal(d[np.argpartition(d, kth, kind=k)][kth], (0, 3, 7, 7))
+
             d = np.array([2, 1])
             d.partition(0, kind=k)
             assert_raises(ValueError, d.partition, 2)
@@ -1551,6 +1557,18 @@ def test_partition_unicode_kind(self):
         assert_raises(ValueError, d.partition, 2, kind=k)
         assert_raises(ValueError, d.argpartition, 2, kind=k)
 
+    def test_partition_fuzz(self):
+        # a few rounds of random data testing
+       for j in range(10, 30):
+           for i in range(1, j - 2):
+               d = np.arange(j)
+               np.random.shuffle(d)
+               d = d % np.random.randint(2, 30)
+               idx = np.random.randint(d.size)
+               kth = [0, idx, i, i + 1]
+               tgt = np.sort(d)[kth]
+               assert_array_equal(np.partition(d, kth)[kth], tgt,
+                                  err_msg="data: %r\n kth: %r" % (d, kth))
 
     def test_flatten(self):
         x0 = np.array([[1, 2, 3], [4, 5, 6]], np.int32)
