Merge pull request #16895 from DevinShanahan/delete-speedup

MAINT: extend delete single value optimization

1 files changed, 12 insertions, 5 deletions

diff --git a/numpy/lib/function_base.py b/numpy/lib/function_base.py
index f4296df2e..ff56196c3 100644
--- a/numpy/lib/function_base.py
+++ b/numpy/lib/function_base.py
@@ -5094,6 +5094,18 @@ def delete(arr, obj, axis=None):
             return new
 
     if isinstance(obj, (int, integer)) and not isinstance(obj, bool):
+        single_value = True
+    else:
+        single_value = False
+        _obj = obj
+        obj = np.asarray(obj)
+        if obj.size == 0 and not isinstance(_obj, np.ndarray):
+            obj = obj.astype(intp)
+        elif obj.size == 1 and not isinstance(_obj, bool):
+            obj = obj.astype(intp).reshape(())
+            single_value = True
+
+    if single_value:
         # optimization for a single value
         if (obj < -N or obj >= N):
             raise IndexError(
@@ -5110,11 +5122,6 @@ def delete(arr, obj, axis=None):
         slobj2[axis] = slice(obj+1, None)
         new[tuple(slobj)] = arr[tuple(slobj2)]
     else:
-        _obj = obj
-        obj = np.asarray(obj)
-        if obj.size == 0 and not isinstance(_obj, np.ndarray):
-            obj = obj.astype(intp)
-
         if obj.dtype == bool:
             if obj.shape != (N,):
                 raise ValueError('boolean array argument obj to delete '