修改2D NumPy數組的每一行中的不同列

Question

我有以下問題：

假設我有一個像這樣定義的數組：

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

我想要做的是利用Numpy多重索引並將幾​​個元素設置為0.為此，我創建了一個向量：

indices_to_remove = [1, 2, 0]

我想要的意思是：

1. 從第一行刪除索引為“1”的元素
2. 從第二行刪除索引為“2”的元素
3. 從第三行刪除索引為“0”的元素

結果應該是數組[[1,0,3],[4,5,0],[0,8,9]]

我設法通過以下代碼獲取我想要修改的元素的值：

values = np.diagonal(np.take(A, indices, axis=1))

但是，這不允許我修改它們。 怎么能解決這個問題？

Answer 1

您可以使用integer array indexing來分配這些零 -

A[np.arange(len(indices_to_remove)), indices_to_remove] = 0

樣品運行 -

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

In [446]: indices_to_remove
Out[446]: [1, 2, 0]

In [447]: A[np.arange(len(indices_to_remove)), indices_to_remove] = 0

In [448]: A
Out[448]: 
array([[1, 0, 3],
       [4, 5, 0],
       [0, 8, 9]])