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.

MSeifert · Accepted Answer · 2016-12-26 12:00:23Z

3

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.

answered Dec 26, 2016 at 12:00

MSeifert

154k41 gold badges356 silver badges377 bronze badges

Sign up to request clarification or add additional context in comments.

3 Comments

mzzx Over a year ago

Just to check: n.ogrid[0:x, 0:y] + [c] adds a third dimension to the ogrid-array, with the elements of c, correct? So this seems very similar to the fancy-indexing answer (which I've accepted as it's even shorter, but this is interesting - thanks!).

Divakar Over a year ago

@mzzx Basically similar answers, as both create broadcasted range arrays to index into the data array. Difference is just in the way those range arrays are created. My post does the broadcasting explicitly, whereas this post does it with the help of a function.

MSeifert Over a year ago

@mzzx Yes, they are equivalent! If you like it short, you could also write it in one line: arr[np.ogrid[0:arr.shape[0], 0:arr.shape[1]] + [c]] :-)

Divakar · Accepted Answer · 2016-12-26 12:30:15Z

1

Here's an approach using fancy-indexing -

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

answered Dec 26, 2016 at 12:30

Divakar

222k19 gold badges273 silver badges374 bronze badges

Comments

Kasravnd · Accepted Answer · 2016-12-26 11:32:20Z

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])

answered Dec 26, 2016 at 11:32

Kasravnd

108k19 gold badges167 silver badges195 bronze badges

2 Comments

hpaulj Over a year ago

Reshape these to (5,2) and the result should look like the mgrid version of the ogrid suggested by another answer.

Kasravnd Over a year ago

@hpaulj Yes,unfortunately I wasn't familiar with that function.

Collectives™ on Stack Overflow

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

3 Answers 3

3 Comments

Comments

2 Comments

Your Answer

Linked

Hot Network Questions

Collectives™ on Stack Overflow

3 Answers 3

3 Comments

Comments

2 Comments

Your Answer

Sign up or log in

Post as a guest

Linked

Related