Slicing a numpy array based on argmax of another numpy array

Question

I have two arrays, and are as follows:

import numpy as np

np.random.seed(42)
a = (np.random.uniform(size=[2, 5, 3]) * 100).astype(int)
b = (np.random.uniform(size=[2, 5, 3]) * 100).astype(int)

Ouput of array a:

array([[[37, 95, 73],
        [59, 15, 15],
        [ 5, 86, 60],
        [70,  2, 96],
        [83, 21, 18]],

       [[18, 30, 52],
        [43, 29, 61],
        [13, 29, 36],
        [45, 78, 19],
        [51, 59,  4]]])

The output of array b is as follows:

array([[[60, 17,  6],
        [94, 96, 80],
        [30,  9, 68],
        [44, 12, 49],
        [ 3, 90, 25]],

       [[66, 31, 52],
        [54, 18, 96],
        [77, 93, 89],
        [59, 92,  8],
        [19,  4, 32]]])

Now I am able to get the argmax of array a using the following code:

idx = np.argmax(a, axis=0)
print(idx)

Output:

array([[0, 0, 0],
       [0, 1, 1],
       [1, 0, 0],
       [0, 1, 0],
       [0, 1, 0]], dtype=int64)

Is it possible to slice array b using the argmax output of array a, so that I get the following output:

array([[60, 17,  6],
       [94, 18, 96],
       [77,  9, 68],
       [44, 92, 49],
       [ 3, 4, 25]])

I tried different ways, but not successful. Kindly help.

Pretty sure there is some dupe but you can do `np.take_along_axis(b, idx[None], axis=0)[0]`. — jdehesa, Jun 15 '20 at 09:59

score 2 · Answer 1 · answered Jun 15 '20 at 10:14

Using numpy advanced indexing:

import numpy as np

np.random.seed(42)
a = (np.random.uniform(size=[2, 5, 3]) * 100).astype(int)
b = (np.random.uniform(size=[2, 5, 3]) * 100).astype(int)

idx = np.argmax(a, axis=0)

_, m, n = a.shape
b[idx, np.arange(m)[:,None], np.arange(n)]
array([[60, 17,  6],
       [94, 18, 96],
       [77,  9, 68],
       [44, 92, 49],
       [ 3,  4, 25]])

Slicing a numpy array based on argmax of another numpy array

1 Answers1