python - 从另一个数组更新 numpy 2d 数组索引

Question

有没有一种不使用循环的 numpy 方法来执行以下操作？

A = np.array([[1,1],[2,2],[3,3],[4,4],[5,5]])
I = np.array([0, 1, 0, 1, 1])
V = np.array([6, 6, 6, 6, 6])

我想更新 A，使其在 I 指定的相应索引中有 6。所以 A 变为 ...

A = np.array([[6,1],[2,6],[6,3],[4,6],[5,6]])

尝试了以下，但他们没有工作..

A[I] = V
A[:,I] = V

标签： pythonarraysnumpy

Answer 1

尝试A[np.arange(len(A)), I] = V：：

In [15]: import numpy as np                                                                                              

In [16]: A = np.array([[1,1],[2,2],[3,3],[4,4],[5,5]])                                                                   

In [17]: I = np.array([0, 1, 0, 1, 1])                                                                                   

In [18]: V = np.array([6, 6, 6, 6, 6])                                                                                   

In [19]: A[np.arange(len(A)), I] = V                                                                                     

In [20]: A                                                                                                               
Out[20]: 
array([[6, 1],
       [2, 6],
       [6, 3],
       [4, 6],
       [5, 6]])