2D Numpy数组的前N个元素，其中N为列表

Question

假设我有一个MxN数组，

>>> 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.]])

，我想获得每个数组的第一个x元素，但x是一个大小为M的列表。因此，如果x是

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

那么我的输出将是

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

我已经尝试使用这种技术和其他类似的技术来做到这一点。这个方法不起作用，因为据我所知，索引输入不能是一个数组。

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

我知道这可以使用for循环来完成，但我希望避免这样做，因为我正在处理大量数据。

此外，输出很可能需要一个数组列表，而不是一个2D数组，因为Numpy似乎不赞成使用不同大小的数组。

Answer 1

你可以通过生成器写你想写的东西但是在循环的最后:

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)))

也许你可以比for循环节省一点(甚至不确定)，但最好是看看这样做的目的是什么，如果你真的关心保存这样的列表，或者你只想要优化的值。