Say I have an M x N array,

>>> M = 5
>>> N = 4
>>> a = np.ones((M, N))

array([[1., 1., 1., 1.],
       [1., 1., 1., 1.],
       [1., 1., 1., 1.],
       [1., 1., 1., 1.],
       [1., 1., 1., 1.]])

and I would like to get the first x elements of each array but where x is a list of size M. So, if x is

array([2, 3, 1, 4])

then my output would be

[[1, 1], 
 [1, 1, 1], 
 [1], 
 [1, 1, 1, 1]]

I have attempted to do this using this technique and others like it. This method does not work because the index input cannot be an array as far as I can tell.

x = [2, 3, 1, 4]
a[:, :x]

I know this could be done using a for-loop but I would like to avoid that since I am working with large amounts of data.

Also, the output will most likely need to be a list of arrays rather than a 2D array since Numpy seems to have deprecated having an array of different-sized arrays.

CodePudding user response：

you could write what you want by generator but at the end of the way at is loop:

import numpy as np
M = 5
N = 4
a = np.ones((M, N))
d = np.array([2, 3, 1, 4])

out = list(map(lambda a:a[0][:a[1]], zip(a, d)))

maybe you could save slightly compared to for loop (not sure even), but best is to see what is the purpose of something like that, if you really care to keep as such list or you only wants the values, to optimise.