From 734ee2c44e4e37745480afa537f1b8af75ea4422 Mon Sep 17 00:00:00 2001 From: DevinShanahan Date: Sat, 18 Jul 2020 07:25:26 -0600 Subject: ENH: extend delete single value optimization Allow arrays of shape (1,) for delete's obj parameter to utilize the optimization for a single value. See #16685. --- numpy/lib/function_base.py | 18 +++++++++++++----- 1 file changed, 13 insertions(+), 5 deletions(-) diff --git a/numpy/lib/function_base.py b/numpy/lib/function_base.py index 1053598b1..99e06d5b8 100644 --- a/numpy/lib/function_base.py +++ b/numpy/lib/function_base.py @@ -4368,6 +4368,19 @@ 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) + + if obj.shape == (1,): + obj = obj.item() + single_value = True + + if single_value: # optimization for a single value if (obj < -N or obj >= N): raise IndexError( @@ -4384,11 +4397,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 ' -- cgit v1.2.1