Find array corresponding to minimal values along an axis in another array

Question

I have two three dimensional arrays, a and b, and want to find the 2D subarray of b with the elements where a had a minimum along the third axis, i.e.

a=n.random.rand(20).reshape((5,2,2))
b=n.arange(20).reshape((5,2,2))
c=n.argmin(a,2) #indices with minimal value of a
d=n.zeros_like(c) #the array I want
for i in range(5):
  for j in range(2):
     d[i,j] = b[i,j,c[i,j]] 

Is there a way I can get these values without the double loop?

I am aware of this answer: replace min value to another in numpy array but if I want this to work for my 3D arrays I'd have to do a lot of reshaping operations - and I'm wondering if there is something simpler.

Answer 1

You can use np.ogrid to create a grid for the other dimensions:

x, y, z = arr.shape  # assuming your array is named "arr"
xyz = np.ogrid[0:x, 0:y] + [c]  # c is your second axis index (the argmin)
arr[xyz]

If it's not the last axis then you can simply use insert because ogrid returns a normal python list containing the indices.