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 ( 3 )
Here is a Numpythonic way:
Here
np.repeat(np.arange(x), y)
will give you the corresponding indices of the first axis.np.tile(np.arange(y), x)
will give you the corresponding indices of the second axis.And for the third one you can just use the flattened shape of
c
.You can use
np.ogrid
to create a grid for the other dimensions:If it's not the last axis then you can simply use
insert
becauseogrid
returns a normal python list containing the indices.Here's an approach using
fancyindexing
