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?
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] ]