algos.kth_smallest

Author: simonjayhawkins

pandas/_libs_numba/algos.py

@@ -230,34 +230,36 @@
 #     return result, counts
 
 
-# @cython.boundscheck(False)
-# @cython.wraparound(False)
-# def kth_smallest(numeric[:] a, Py_ssize_t k) -> numeric:
-#     cdef:
-#         Py_ssize_t i, j, l, m, n = a.shape[0]
-#         numeric x
-
-#     with nogil:
-#         l = 0
-#         m = n - 1
-
-#         while l < m:
-#             x = a[k]
-#             i = l
-#             j = m
-
-#             while 1:
-#                 while a[i] < x: i += 1
-#                 while x < a[j]: j -= 1
-#                 if i <= j:
-#                     swap(&a[i], &a[j])
-#                     i += 1; j -= 1
-
-#                 if i > j: break
-
-#             if j < k: l = i
-#             if k < i: m = j
-#     return a[k]
+@numba.njit
+def kth_smallest(a: np.ndarray, k):
+    n = a.shape[0]
+
+    l = 0
+    m = n - 1
+
+    while l < m:
+        x = a[k]
+        i = l
+        j = m
+
+        while 1:
+            while a[i] < x:
+                i += 1
+            while x < a[j]:
+                j -= 1
+            if i <= j:
+                a[i], a[j] = a[j], a[i]
+                i += 1
+                j -= 1
+
+            if i > j:
+                break
+
+        if j < k:
+            l = i
+        if k < i:
+            m = j
+    return a[k]
 
 
 # # ----------------------------------------------------------------------

pandas/core/algorithms.py

@@ -28,6 +28,7 @@
     iNaT,
     lib,
 )
+from pandas._libs_numba import algos as algos_numba
 from pandas._typing import (
     AnyArrayLike,
     ArrayLike,
@@ -1290,7 +1291,7 @@ def compute(self, method: str) -> Series:
         narr = len(arr)
         n = min(n, narr)
 
-        kth_val = algos.kth_smallest(arr.copy(), n - 1)
+        kth_val = algos_numba.kth_smallest(arr.copy(), n - 1)
         (ns,) = np.nonzero(arr <= kth_val)
         inds = ns[arr[ns].argsort(kind="mergesort")]
 

setup.cfg

@@ -76,6 +76,7 @@ ignore =
     W504,  # line break after binary operator
     E402,  # module level import not at top of file
     E731,  # do not assign a lambda expression, use a def
+    E741,  # ambiguous variable name
     S001   # found modulo formatter (incorrect picks up mod operations)
 exclude =
     doc/sphinxext/*.py,