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

时间:2018-06-27 18:07:56

标签: python numpy vectorization

让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将其向量化为循环。

1 个答案:

答案 0 :(得分: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