如何以矢量化方式填充 NumPy 数组列表？

Question

我正在尝试找到一种矢量化方式（或者至少比使用循环更好）来从 2D NumPy 数组列表中创建一个 3D NumPy 数组。现在，我有一个列表 L，它看起来像：

L = [ np.array([[1,2,3], [4,5,6]]), np.array([[8,9,10]]), ...]

每个 NumPy 数组的第二维大小相同（在上面的例子中，它是 3）。但是第一个维度有不同的大小。

我的目标是创建一个包含上述数据的 3D NumPy 数组 M。我一直在尝试使用 np.pad() 函数，因为我的每个数组的第一个维度都有一个最大大小，但看起来它只会对列表中的单个元素进行操作。然后我可以使用该函数做我想做的事并遍历每个数组。但是，如果可能，我想使用矢量化方法在没有循环的情况下执行此操作。有什么技术可以做到这一点吗？

这个问题与this one 相关，但我希望一次在我的整个列表中进行此操作。

Answer 1

首先让我们看一下将一维数组填充到常见大小的常见任务。

In [441]: alist = [np.ones((2,),int),np.zeros((1,),int)+2, np.zeros((3,),int)+3]
In [442]: alist
Out[442]: [array([1, 1]), array([2]), array([3, 3, 3])]

明显的迭代方法：

In [443]: [np.hstack((arr, np.zeros((3-arr.shape[0]),int))) for arr in alist]
Out[443]: [array([1, 1, 0]), array([2, 0, 0]), array([3, 3, 3])]
In [444]: np.stack(_)
Out[444]: 
array([[1, 1, 0],
       [2, 0, 0],
       [3, 3, 3]])

一个聪明的选择。它仍然需要迭代来确定大小，但剩下的就是整个数组“矢量化”：

In [445]: sizes = [arr.shape[0] for arr in alist]
In [446]: sizes
Out[446]: [2, 1, 3]

使用填充值制作输出数组：

In [448]: res = np.zeros((3,3),int)

制作一个聪明的面具（@Divakar 最先提出这个）

In [449]: np.array(sizes)[:,None]>np.arange(3)
Out[449]: 
array([[ True,  True, False],
       [ True, False, False],
       [ True,  True,  True]])

然后将“扁平化”输入映射到 res：

In [450]: res[_]=np.hstack(alist)
In [451]: res
Out[451]: 
array([[1, 1, 0],
       [2, 0, 0],
       [3, 3, 3]])

我认为这个过程可以扩展到您的 2d=>3d 案例。但这需要一些工作。我尝试直接这样做，发现我在涂抹面膜时迷失了方向。这就是为什么我决定首先布局 1d=>2d 案例。有足够的跳出框框思考，我每次都必须重新制定细节。

2d=>3d

In [457]: a2list = [np.ones((2,3),int),np.zeros((1,3),int)+2, np.zeros((3,3),int)+3]
In [458]: [np.vstack((arr, np.zeros((3-arr.shape[0],arr.shape[1]),int))) for arr in a2list]
Out[458]: 
[array([[1, 1, 1],
        [1, 1, 1],
        [0, 0, 0]]),
 array([[2, 2, 2],
        [0, 0, 0],
        [0, 0, 0]]),
 array([[3, 3, 3],
        [3, 3, 3],
        [3, 3, 3]])]
In [459]: np.stack(_)
Out[459]: 
array([[[1, 1, 1],
        [1, 1, 1],
        [0, 0, 0]],

       [[2, 2, 2],
        [0, 0, 0],
        [0, 0, 0]],

       [[3, 3, 3],
        [3, 3, 3],
        [3, 3, 3]]])

现在介绍“矢量化”方法：

In [460]: sizes = [arr.shape[0] for arr in a2list]
In [461]: sizes
Out[461]: [2, 1, 3]
In [462]: np.array(sizes)[:,None]>np.arange(3)
Out[462]: 
array([[ True,  True, False],
       [ True, False, False],
       [ True,  True,  True]])
In [463]: res = np.zeros((3,3,3),int)

以及掩码中的相应索引：

In [464]: I,J=np.nonzero(Out[462])
In [465]: I
Out[465]: array([0, 0, 1, 2, 2, 2])
In [466]: J
Out[466]: array([0, 1, 0, 0, 1, 2])
In [467]: res[I,J,:] = np.vstack(a2list)
In [468]: res
Out[468]: 
array([[[1, 1, 1],
        [1, 1, 1],
        [0, 0, 0]],

       [[2, 2, 2],
        [0, 0, 0],
        [0, 0, 0]],

       [[3, 3, 3],
        [3, 3, 3],
        [3, 3, 3]]])