根据长度列表从numpy数组中选择元素

Question

假设我有一个像这样的numpy数组：

arr = np.array([[1,1,1,1], [2,2,2,2], [3,3,3,3], [4,4,4,4])

我也有一个确定预期长度的列表：

lens = [1,2,3,4] 

是否有一种优雅而Pythonic的方式来返回一个新数组，并使用该lens变量选择了每个对应的元素？

输出应为：

[[1], [2,2],[3,3,3], [4,4,4,4]]

Answer 1

如果每个Numpy arr列表元素的大小小于5，则可以使用。

用途zip：

[a[:i].tolist() for a, i in zip(arr, lens)]

输出：

[[1], [2, 2], [3, 3, 3], [4, 4, 4, 4]]