Numpy 按特定数量的元素将数组拆分为子数组

Question

例如

我有一个积分列表

a = np.array([[0,0,0],
              [1,1,1],
              [2,2,2],
              [3,3,3],
              [4,4,4],
              [5,5,5],
              [6,6,6],
              [7,7,7],
              [8,8,8],
              [9,9,9]])

我还有另一个数组表示元素的个数

b = np.array([2,0,3,5])

如何根据数组b的元素个数拆分数组a得到输出

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

Answer 1

您可以在b上使用numpy.split using cumsum来获得分割点：

out = np.split(a, b.cumsum()[:-1])

输出：

[array([[0, 0, 0],
        [1, 1, 1]]),
 array([], shape=(0, 3), dtype=int64),
 array([[2, 2, 2],
        [3, 3, 3],
        [4, 4, 4]]),
 array([[5, 5, 5],
        [6, 6, 6],
        [7, 7, 7],
        [8, 8, 8],
        [9, 9, 9]])]

如果你想要列表：

out = [x.tolist() for x in np.split(a, b.cumsum()[:-1])]

输出：

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

中级：

b.cumsum()[:-1]
# array([2, 2, 5])

Numpy 按特定数量的元素将数组拆分为子数组

Numpy Split array to sub arrays by specific number of elements

python

arrays

numpy

split