时间戳

首页 > 解决方案 > 在不使用循环的情况下，在 numpy 中选择具有 (n,n) 形状掩码的 (n,n,2) numpy 数组的元素

python - 在不使用循环的情况下，在 numpy 中选择具有 (n,n) 形状掩码的 (n,n,2) numpy 数组的元素

问题描述

我有一个 (n,n,2) numpy 数组，我想根据 (n,n) 掩码选择其元素，而不使用循环。有没有办法在 numpy 中对这个操作进行矢量化？说我有一个 numpy 数组

X = array([[[18,  8],
        [ 9,  2],
        [11,  4],
        [18, 14]],

       [[ 8, 10],
        [13,  5],
        [13,  6],
        [13, 18]],

       [[ 8,  4],
        [ 2, 13],
        [19, 11],
        [ 3, 15]],

       [[12,  6],
        [ 7,  3],
        [19, 17],
        [ 1, 12]]])

和一个面具

M = array([[1, 0, 0, 0],
           [1, 1, 0, 0],
           [0, 0, 1, 0],
           [0, 0, 0, 0]])

将 X 中的每个二维条目视为一个元素，有没有办法使用掩码 M 来选择 X 的元素？也就是说，如果掩码 M 中的对应元素为 1，则选择 X 中的二维元素。

所以上面的例子会返回

[
 [[18,  8]],
 [[ 8, 10],
  [13,  5]],
 [[19, 11]],
 []
]

标签： pythonarraysnumpy

解决方案

In [393]: I,J = np.nonzero(M)                                                                  
In [394]: I,J                                                                                  
Out[394]: (array([0, 1, 1, 2]), array([0, 0, 1, 2]))
In [395]: X[I,J,:]                                                                             
Out[395]: 
array([[18,  8],
       [ 8, 10],
       [13,  5],
       [19, 11]])

有人指出，您不仅想要这些元素，还希望它们按行分组。为此，列表推导做得很好：

In [407]: [x[m.astype(bool)] for x,m in zip(X,M)]                                              
Out[407]: 
[array([[18,  8]]), array([[ 8, 10],
        [13,  5]]), array([[19, 11]]), array([], shape=(0, 2), dtype=int64)]
In [408]: [x[m.astype(bool)].tolist() for x,m in zip(X,M)]                                     
Out[408]: [[[18, 8]], [[8, 10], [13, 5]], [[19, 11]], []]

是的，它使用了一个循环，但这就是您在获取可变大小项目列表时所期望的。

vectorized, 无循环解决方案仅适用于常规数组结果。

推荐阅读