# [Numpy-discussion] using numpy.argmax to index into another array

Moroney, Catherine M (388D) Catherine.M.Moroney@jpl.nasa....
Wed Oct 31 18:23:33 CDT 2012

```Hello Everybody,

I have the following problem that I would be interested in finding an easy/elegant solution to.
I've got it working, but my solution is exceedingly clunky and I'm sure that there must be a
better way.

Here is the (boiled-down) problem:

I have 2 different 3-d array, shaped (4,4,4), "a" and "b"

I can find the max values and location of those max values in "a" in the 0-th dimension
using max and argmax resulting in a 4x4 matrix.  So far, very easy.

I then want to find the values in "b" that correspond to the maximum values in a.
This is where I got stuck.

Below find the sample code I used (pretty clumsy stuff ...)
Can somebody show a better (less clumsy) way of finding bmax
in the code excerpt below?

>>> a = numpy.random.randint(0, high=15, size=(4,4,4))
>>> b = numpy.random.randint(0, high=15, size=(4,4,4))
>>> amax = a.argmax(axis=0)

>>> idx2 = [0, 1, 2, 3, 0, 1, 2, 3, 0, 1, 2, 3, 0, 1, 2, 3]
>>> idx1 = [0, 0, 0, 0, 1, 1, 1, 1, 2, 2, 2, 2, 3, 3, 3, 3]
>>> idx_all = (amax.flatten(), numpy.array(idx1), numpy.array(idx2))

>>> bmax = b[idx_all].reshape(4,4)

>>> a
array([[[ 7,  1,  7, 10],
[11,  5,  9,  0],
[13,  1, 13, 14],
[ 4, 13,  4,  7]],

[[ 2, 11,  4,  5],
[ 6, 12, 14,  9],
[ 0,  4,  5, 14],
[13,  5,  3,  4]],

[[ 1,  4,  5,  4],
[ 7, 13,  1, 11],
[ 9,  2, 10,  4],
[ 3,  4, 10, 13]],

[[14,  6,  7,  1],
[ 8,  0,  8,  2],
[ 6,  8,  7, 12],
[ 4,  9,  0, 12]]])
>>> idx
array([[3, 1, 0, 0],
[0, 2, 1, 2],
[0, 3, 0, 0],
[1, 0, 2, 2]])
>>> b
array([[[ 4,  2,  3,  2],
[ 9,  6,  1,  4],
[ 1, 10, 13, 11],
[10,  8,  1, 10]],

[[ 3, 13,  7,  7],
[ 4,  7,  9,  7],
[ 3, 13, 10, 10],
[13,  8, 14,  6]],

[[14, 10, 13,  6],
[13,  2, 12, 12],
[ 9,  6,  3,  8],
[ 0,  7,  8, 11]],

[[12, 10,  3,  0],
[11,  7,  0,  4],
[ 9,  7, 11, 12],
[ 7, 12,  1,  1]]])
>>> bmax = b[idx_all].reshape(4,4)
>>> bmax
array([[12, 13,  3,  2],
[ 9,  2,  9, 12],
[ 1,  7, 13, 11],
[13,  8,  8, 11]])

Catherine

```