## 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.

Show source

## Answers to Find array corresponding to minimal values along an axis in another array ( 3 )

1. Here is a Numpythonic way:

``````In [83]: x, y, z = a.shape
In [84]: b[np.repeat(np.arange(x), y), np.tile(np.arange(y), x), c.ravel()].reshape(x, y)
``````

Here `np.repeat(np.arange(x), y)` will give you the corresponding indices of the first axis.

``````In [86]: np.repeat(np.arange(x), y)
Out[86]: array([0, 0, 1, 1, 2, 2, 3, 3, 4, 4])
``````

`np.tile(np.arange(y), x)` will give you the corresponding indices of the second axis.

``````In [87]: np.tile(np.arange(y), x)
Out[87]: array([0, 1, 0, 1, 0, 1, 0, 1, 0, 1])
``````

And for the third one you can just use the flattened shape of `c`.

``````In [88]: c.ravel()
Out[88]: array([1, 0, 1, 1, 1, 0, 1, 1, 0, 0])
``````
2. 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.

3. Here's an approach using `fancy-indexing` -

``````m,n,r = b.shape
d_out = b[np.arange(m)[:,None],np.arange(n),c]
``````