python - Set value in 2D Numpy array based on row sum

Question

Is this possible to accomplish with Numpy and with good performance?

Initial 2D array:

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

If the sum of each row is less than 4, set the last item in each row to 1:

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

Divide each item in each row with the sum of each row and get this result:

array([[0, 0.25, 0.25, 0.25, 0.25, 0],
   [0, 0, 0.5, 0, 0, 0.5],
   [0.5, 0, 0, 0, 0, 0.5]])

标签： pythonarraysnumpy

Answer 1

numpy.where在此处查找与您的条件匹配的行也很有用：

import numpy as np
a = np.array([[0, 1, 1, 1, 1, 0],
              [0, 0, 1, 0, 0, 0],
              [1, 0, 0, 0, 0, 1]])

a[np.sum(a,axis=1) < 4, -1] = 1
a = a/a.sum(axis=1)[:,None]

print(a)

# Output 
# [[0.   0.25 0.25 0.25 0.25 0.  ]
#  [0.   0.   0.5  0.   0.   0.5 ]
#  [0.5  0.   0.   0.   0.   0.5 ]]

PS：在@tel 建议后编辑:)