CODEWITHSUNDEEP
pythonarraysnumpy
cs95
cs95

Reputation: 402483

2D argmax on a numpy array

Starting with a numpy array v:

v = \
np.array([[0, 0, 0, 0, 0],
          [0, 0, 1, 0, 0],
          [1, 0, 0, 0, 0],
          [0, 1, 0, 0, 0],
          [0, 0, 1, 0, 1],
          [0, 0, 0, 0, 1],
          [0, 0, 0, 1, 1]])

I want to find the first non-zero value in each column and multiply those values by 100.

My desired result is:

array([[  0,   0,   0,   0,   0],
       [  0,   0, 100,   0,   0],
       [100,   0,   0,   0,   0],
       [  0, 100,   0,   0,   0],
       [  0,   0,   1,   0, 100],
       [  0,   0,   0,   0,   1],
       [  0,   0,   0, 100,   1]])

I thought to approach this problem by taking the argmax along each axis:

i = v.argmax(0)
j = v.argmax(1)
v[i, j] *= 100

I know I'm not using i and j correctly, so how can I fix this?

Upvotes: 1

Views: 504

Answers (1)

Sebastian Mendez
Sebastian Mendez

Reputation: 2981

You just want to range through the columns

v[v.argmax(axis=0), np.arange(v.shape[1])] *= 100

Upvotes: 3

Related Questions