How to get mask from numpy array according to all values on third axis
<p>I have numpy array like:</p>

<pre><code>A = np.zeros((X,Y,Z)) 
</code></pre>

<p>Later I fill the array with values and I need to find x,y coordinates according to values in arrays on z axis to replace it with different array.</p>

<p>Example with<code>Z=2</code> I have array</p>

<pre><code>A = [ [ [1,2], [2,1], [0, 0] ],
      [ [1,1], [1,1], [0, 0] ], ]
</code></pre>

<p>I need something like</p>

<pre><code>A[ where A[x,y,:] == [1,2] ] = [2,1]
</code></pre>

<p>What will produce array</p>

<pre><code>A = [ [ [2,1], [2,1], [0, 0] ],
      [ [1,1], [1,1], [0, 0] ], ]
</code></pre>

<p>Is it possible to achieve that somehow simple? I would like to avoid iteration over the x,y coordinates.</p>

<p>One vectorized approach -</p>

<pre><code>A[(A == [1,2]).all(-1)] = [2,1]
</code></pre>

<p>Sample run -</p>

<pre><code>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]]])
</code></pre>

python arrays numpy vectorization