Векторизованное обновление массива с использованием других элементов массива в качестве индекса

Question

Пусть A, C и B - массивы с одинаковым количеством строк. Я хочу обновить 0-й элемент A [0], 2-й элемент A [1] и т. Д. То есть обновить B [i] -й элемент A [i] до C [i]

import numpy as np
A = np.array([[1,2,3],[3,4,5],[5,6,7],[0,8,9],[3,7,5]])
B = np.array([0,2,1,2,0])
C = np.array([8,9,6,5,4])
for i in range(5):
A[i, B[i]] = C[i]
print ("FOR", A)
A = np.array([[1,2,3],[3,4,5],[5,6,7],[0,8,9],[3,7,5]])
A[:,B[:]] = C[:]
print ("Vectorized, A", A)

Выход:

FOR [[8 2 3]
[3 4 9]
[5 6 7]
[0 8 5]
[4 7 5]]
Vectorized, A [[4 6 5]
[4 6 5]
[4 6 5]
[4 6 5]
[4 6 5]]

Цикл for и векторизация дали разные результаты. Я не уверен, как векторизовать это для цикла с помощью Numpy.

Answer 1

Причина, по которой ваш подход не работает, заключается в том, что вы передаете целое B в качестве индекса столбца и заменяете их C, вместо этого вам нужно указать индекс строки и индекс столбца. Поскольку вы просто хотите изменить первые 4 строки, вы можете просто использовать np.arange(4), чтобы выбрать строки B[:4] столбцы и C[:4] элементы замены.

In [26]: A[np.arange(4),B[:4]] = C[:4]

In [27]: A
Out[27]: 
array([[8, 2, 3],
       [3, 4, 9],
       [5, 6, 7],
       [0, 8, 5],
       [3, 7, 5]])

Обратите внимание, что если вы хотите обновить весь массив, как упомянуто в комментариях @Warren, вы можете использовать следующий подход:

A[np.arange(A.shape[0]), B] = C