如何根据第三轴上的所有值从 numpy 数组中获取掩码答案

【问题标题】：How to get mask from numpy array according to all values on third axis如何根据第三轴上的所有值从 numpy 数组中获取掩码
【发布时间】：2016-10-08 10:40:41
【问题描述】：

我有像这样的 numpy 数组：

A = np.zeros((X,Y,Z))

稍后我用值填充数组，我需要根据z轴上数组中的值找到x，y坐标以将其替换为不同的数组。

以Z=2 为例我有数组

A = [ [ [1,2], [2,1], [0, 0] ],
      [ [1,1], [1,1], [0, 0] ], ]

我需要类似的东西

A[ where A[x,y,:] == [1,2] ] = [2,1]

什么会产生数组

A = [ [ [2,1], [2,1], [0, 0] ],
      [ [1,1], [1,1], [0, 0] ], ]

有可能以某种方式简单地实现吗？我想避免在 x,y 坐标上进行迭代。

【问题讨论】：

您能添加一个示例案例和预期的输出吗？
@Divakar 完成。如果还不够，请告诉我。

标签： python arrays numpy vectorization

【解决方案1】：

一种矢量化方法 -

A[(A == [1,2]).all(-1)] = [2,1]

示例运行 -

In [15]: A
Out[15]: 
array([[[1, 2],
        [2, 1],
        [0, 0]],

       [[1, 1],
        [1, 2],
        [0, 0]]])

In [16]: A[(A == [1,2]).all(-1)] = [2,1]

In [17]: A
Out[17]: 
array([[[2, 1],
        [2, 1],
        [0, 0]],

       [[1, 1],
        [2, 1],
        [0, 0]]])

【讨论】：