CODEWITHSUNDEEP
pythonarraysnumpyarray-broadcasting
tczj
tczj

Reputation: 448

How to effectively slice numpy array

I want to implement the following code in a matrix-operation way instead of with a for loop.

a = np.random.randint(0, 7, (4,3))
b = np.random.randint(0, 6, (4,3,2))
c = None
for idx in xrange(a.shape[0]):
     max_idx = np.argmax(a[idx])
     ex_b = b[idx, max_idx].reshape(1, -1)
     if c is None:
         c = ex_b
     else:
         c = np.concatenate((c, ex_b), axis=0)

Basically, I want to first get the indices of maximum value in a's second dimension. And then I want to extract the corresponding third dimension values in b according to these indices.

For example:

a:
array([[5, 4, 1],
       [3, 1, 3],
       [4, 1, 2],
       [0, 0, 5]])
b:
array([[[1, 3], [1, 4], [5, 0]],
       [[2, 4], [2, 2], [1, 2]],
       [[2, 1], [1, 2], [4, 5]],
       [[4, 0], [5, 5], [0, 2]]])

then np.argmax(a, axis=1) will give array([0, 0, 0, 2]) So c[0] = b[0][0], c[1]=b[1]b[0], c[2]=b[2][0], c[3]=b[3][2]

I think this for loop will slow down the running speed, is there any more elegant way to achieve this in a faster way?

Upvotes: 2

Views: 194

Answers (1)

Paul Panzer
Paul Panzer

Reputation: 53029

You can use fancy indexing like so:

b[np.arange(4), np.argmax(a, axis=-1)]
# array([[1, 3],
#        [2, 4],
#        [2, 1],
#        [0, 2]])

Upvotes: 2

Related Questions