矢量化更新 numpy 数组，使用另一个 numpy 数组元素作为索引

Question

设 A，C 和 B 是具有相同行数的 numpy 数组。 我想更新 A[0] 的第 0 个元素、A[1] 的第 2 个元素等。也就是说，将 A[i] 的第 B[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