Merge pull request numpy#4837 from juliantaylor/select-bug

charris · charris · commit 7a2b14ae0adb · 2014-07-04T13:22:34.000-06:00
BUG: wrong selection for orders falling into equal ranges
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)
