numpyindexingadvanced-indexing

NumPy, advanced indexing and slicilng


I am trying to figure out how to index both "standalone" elements and ranges simultaneously in NumPy.

Suppose we have an array

arr = np.arange(0, 20)

and we want to take, say, 2nd, 4th, 6th elements and also all elements from 10th to the end of the array. What I coded to achieve that was

arr[np.hstack([2, 4, 6, np.arange(10, len(arr))])]

but this doesn't look very elegant. What I would like is something like that:

arr[[2, 4, 6, 10:]]

The code apparently doesn't work, but is there anything as simple as this?


Solution

  • Unlucky for you, numpy requires a list of only integers (not slices) for advanced indexing purposes.

    Lucky for you, the numpy devs foresaw your use case and provided an index trick to transform such a mixed index to an array of indices: np.r_

    >>> import numpy as np
    
    >>> np.r_[2, 4, 6, 10:20]
    array([ 2,  4,  6, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19])
    

    Note that you need to specify an end for the slice as np.r_ has no idea what the length of your array arr is.

    So you can simply index like:

    arr[ np.r_[2, 4, 6, 10:20] ]