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

Question

我有 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 坐标进行迭代。

Answer 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]]])