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

【问题标题】：Vectorized Update numpy array using another numpy array elements as index使用另一个 numpy 数组元素作为索引的矢量化更新 numpy 数组
【发布时间】：2018-06-28 02:48:08
【问题描述】：

让 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 对这个 for 循环进行矢量化。

【问题讨论】：

A[np.arange(A.shape[0]), B] = C
@WarrenWeckesser 这也会改变最后一行。
文本中的描述听起来像最后一行应该更新。我认为range(4) 是一个印刷错误。
@WarrenWeckesser 是的，这就是我一开始的假设并犹豫是否发布答案，但正如我在回答中解释的那样，这样做是完全可能的。
Surabhi，C的最后一行不应该是[4, 7, 5]吗？

标签： python numpy vectorization

【解决方案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 在 cmets 中提到的，您可以使用以下方法：

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

【讨论】：