选择numpy中具有(n，n)形状掩码的(n，2)numpy数组的元素，而不使用循环

Question

我有一个(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的元素？也就是说，如果X中的2-D元素在掩模M中的对应元素为1，则选择该2-D元素。

所以上面的例子将返回

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

Answer 1

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，无循环解决方案仅适用于规则数组结果。

Answer 2

您可以使用：print(X[M==1](