How to add values to only certain index in Numpy 2D Matrix-CodePudding

I have a 2d Matrix

matrix = np.array([[1,2],[3,4],[5,6]])
index = np.array([0, 1, 1])
add_value = np.array([1, 2, 3])

I want to add add_value to matrix but only to the elements corresponding to index in the index list. For example, 1 in add_value should be added to the first element in [1,2], which is 1, resulting in 2. So the output should be

np.array([[2,2],[3,6],[5,9]])

CodePudding user response：

Use a simple multi-dimensional indexing:

matrix[np.arange(matrix.shape[0]), index]  = add_value

Or using python builtins:

matrix[tuple(zip(*enumerate(index)))]  = add_value

Output:

array([[2, 2],
       [3, 6],
       [5, 9]])

CodePudding user response：

for i, x in enumerate(add_value):
    matrix[i][index[i]]  = x