二维 Numpy 数组的前 N 个元素，其中 N 是一个列表

Question

假设我有一个M x N数组，

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

那么我的 output 就是

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

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

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

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

此外，output 很可能需要是 arrays 的列表而不是二维数组，因为 Numpy 似乎已弃用不同大小的数组 arrays。

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 循环相比，也许你可以稍微节省一点（甚至不确定），但最好是看看这样的目的是什么，如果你真的关心保留这样的列表或者你只想要值，优化。