python - 如何使用多个索引多次拆分 numpy 数组？

Question

我想在不同的索引处多次拆分数组。

例子：

y = [0, 1, 2, 3, 4, 5, 6, 7, 8, 9]

我想在索引 2、4、7 处拆分此数组以生成以下结果：

res = [[[0, 1], [2, 3, 4, 5, 6, 7, 8, 9]],
       [[0, 1, 2, 3], [4, 5, 6, 7, 8, 9]],
       [[0, 1, 2, 3, 4, 5, 6], [7, 8, 9]]]

我尝试使用 numpy.split 但它似乎只能拆分一次，而不是多次拆分以生成结果矩阵。有没有办法使用 numpy 的矢量化操作而不使用循环来做到这一点？

标签： pythonnumpysplit

Answer 1

我认为没有矢量化运算符。结果本身不是典型的矩阵，因为它不是矩形的。或者，您可以简单地使用列表推导。

y = [0, 1, 2, 3, 4, 5, 6, 7, 8, 9]

split_at = [2, 4, 7]
output = [[y[:x], y[x:]] for x in split_at]
print(output) # [[[0, 1], [2, 3, 4, 5, 6, 7, 8, 9]], [[0, 1, 2, 3], [4, 5, 6, 7, 8, 9]], [[0, 1, 2, 3, 4, 5, 6], [7, 8, 9]]]