using numpy.argmax to index into another array
by Moroney, Catherine M (388D)
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
