I have a 3 dimensional numpy array with dimensions (x = 2, y = 2, z = 3), as shown below
a = [[[0,1,2],[3,4,5]],
[[6,7,8],[9,10,11]]]
I want to get first N elements from each (x, y) element of a, where N is defined in another array of size x,y. E.g.
b = [[1,2],
[0,0]]
The result would should be
c = [[[0,1],[3,4,5]],
[[6],[9]]
How can I do this without loops?
You can create a mask over the elements in a that are within the given ranges and use this mask whenever you need to operate on a. If you want the encapsulate the mask and array in one entity, have a look at the numpy.ma module.
a = np.array([[[0,1,2],[3,4,5]],
[[6,7,8],[9,10,11]]])
b = np.array([[1,2],
[0,0]])
mask = np.arange(3)[None,None,:] <= b[:,:,None]
a[mask]
Output:
array([0, 1, 3, 4, 5, 6, 9])
But if you want the output as a non-homogeneous array, I cannot think of a better way to do that than using loops.
