ENH: add max fastpath to partition for nan detection

juliantaylor · juliantaylor · commit a676f4eb6c09 · 2014-02-22T11:09:07.000+01:00
Allows low overhead check for NaN via isnan(partition(d, (x, -1))[-1])
diff --git a/numpy/core/src/npysort/selection.c.src b/numpy/core/src/npysort/selection.c.src
@@ -69,6 +69,7 @@ static NPY_INLINE void store_pivot(npy_intp pivot, npy_intp kth,
  *         npy_uint, npy_long, npy_ulong, npy_longlong, npy_ulonglong,
  *         npy_ushort, npy_float, npy_double, npy_longdouble, npy_cfloat,
  *         npy_cdouble, npy_clongdouble#
+ * #inexact = 0*11, 1*7#
  */
 
 static npy_intp
@@ -322,6 +323,20 @@ int
         store_pivot(kth, kth, pivots, npiv);
         return 0;
     }
+    else if (@inexact@ && kth == num - 1) {
+        /* useful to check if NaN present via partition(d, (x, -1)) */
+        npy_intp k;
+        npy_intp maxidx = low;
+        @type@ maxval = v[IDX(low)];
+        for (k = low + 1; k < num; k++) {
+            if (!@TYPE@_LT(v[IDX(k)], maxval)) {
+                maxidx = k;
+                maxval = v[IDX(k)];
+            }
+        }
+        SWAP(SORTEE(kth), SORTEE(maxidx));
+        return 0;
+    }
 
     /* dumb integer msb, float npy_log2 too slow for small parititions */
     {
diff --git a/numpy/core/tests/test_multiarray.py b/numpy/core/tests/test_multiarray.py
@@ -1299,6 +1299,16 @@ def test_partition(self):
             mid = x.size // 2 + 1
             assert_equal(np.partition(x, mid)[mid], mid)
 
+            # max
+            d = np.ones(10); d[1] = 4;
+            assert_equal(np.partition(d, (2, -1))[-1], 4)
+            assert_equal(np.partition(d, (2, -1))[2], 1)
+            assert_equal(d[np.argpartition(d, (2, -1))][-1], 4)
+            assert_equal(d[np.argpartition(d, (2, -1))][2], 1)
+            d[1] = np.nan
+            assert_(np.isnan(d[np.argpartition(d, (2, -1))][-1]))
+            assert_(np.isnan(np.partition(d, (2, -1))[-1]))
+
             # equal elements
             d = np.arange((47)) % 7
             tgt = np.sort(np.arange((47)) % 7)
