In NumPy, how to extract a range from a 1d array without using np.s_?

Viewed 252

Assume an array like so:

a = np.arange(10)

I'd like to delete the numbers from index 2 to 5.

I can do it like this:

a = np.delete(a, np.s_[2:6])

Now a contains [0, 1, 6, 7, 8, 9]. However this function is not supported by Numba, and I need to compile this code using Numba.

I would need to accomplish the same using only "basic" NumPy functions (anything here is OK: https://numba.pydata.org/numba-doc/dev/reference/numpysupported.html). Unfortunately the s_ object is not supported.

How can I accomplish this? Its OK if I need to make more than one call or tmp arrays.

3 Answers

The temporary arrays will be created no matter how you do it. You can use some very simple indexing to get what you want:

a = np.arange(10)
a = np.delete(a, slice(2, 6))

The documentation for s_ pretty much tells you how to do this in the notes. A 1D call to s_ is mostly just shorthand for slice.

Using delete is probably the right choice here because it will allocate the output more efficiently than manually slicing the beginning and end and concentrating.

I would advise against deleting from array in numpy as it can be slow (specially for longer arrays, since it copies). Using masks is another way (not sure if numba supports it, worth a try). If you can try to do your operations on masked array:

b = a[2:6]
#[2 3 4 5]


#Try to do operations on masked array
a = np.ma.array(a, mask=False)
a.mask[2:6] = True
#[0 1 -- -- -- -- 6 7 8 9]

#if you insist on deleting masked elements
a = a.compressed()
#[0 1 6 7 8 9]

One way using numpy.arange:

from numba import njit

@njit
def nb_delete(arr, i, j):
    return np.delete(arr, np.arange(i, j))

nb_delete(np.arange(10), 2,6)

Output:

array([0, 1, 6, 7, 8, 9])
Related