Assign to any column in each row with np.array

>>> import numpy as np
>>> A = np.zeros((4,2))
>>> A
array([[ 0.,  0.],
       [ 0.,  0.],
       [ 0.,  0.],
       [ 0.,  0.]])
>>> index = np.array([0,1,0,1])
>>> index
array([0, 1, 0, 1])
>>> values = np.array([10,20,10,10])
>>> values
array([10, 20, 10, 10])
>>> A[np.arange(4),index] = values
>>> A
array([[ 10.,   0.],
       [  0.,  20.],
       [ 10.,   0.],
       [  0.,  10.]])

Recommended Posts