title: In a 3d array in numpy. How can I extract indices of max element in third dimension?
body: <p>Example input 3D array of shape (2,2,2): </p>
  
  <pre><code>[[[ 1, 2],
    [ 4, 3]],
   [[ 5, 6],
    [ 8, 7]]]
  </code></pre>
  
  <p>My 3d array has a shape of (N, N, N), in above example N = 2.</p>
  
  <p>I need to get all indices such that index for third dimension belongs to max element in third dimension, Output for above 3D array: </p>
  
  <pre><code>[[0, 0, 1],  # for element 2
   [0, 1, 0],  # for element 4
   [1, 0, 1],  # for element 6
   [1, 1, 0]]  # for element 8
  </code></pre>
  
  <p>It would be great if I can do that with <code>argmax</code> or <code>argwhere</code> function. I want to avoid iteration and see if its possible to do this using numpy functions. </p>