Numpy: convert index in one dimension into many dimensions

Question

Many array methods return a single index despite the fact that the array is multidimensional. For example:

a = rand(2,3)
z = a.argmax()

For two dimensions, it is easy to find the matrix indices of the maximum element:

a[z/3, z%3]

But for more dimensions, it can become annoying. Does Numpy/Scipy have a simple way of returning the indices in multiple dimensions given an index in one (collapsed) dimension? Thanks.

Answer 1

Got it!

a = X.argmax()
(i,j) = unravel_index(a, X.shape)

Answer 2

I don't know of an built-in function that does what you want, but where this has come up for me, I realized that what I really wanted to do was this:

given 2 arrays a,b with the same shape, find the element of b which is in the same position (same [i,j,k...] position) as the maximum element of a

For this, the quick numpy-ish solution is:

j = a.flatten().argmax()
corresponding_b_element = b.flatten()[j]

Vince Marchetti