[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






More information about the NumPy-Discussion mailing list