python - 如何在没有for循环的情况下通过多个索引获取/设置numpy矩阵中的元素？

Question

例如，假设我们有

a = zeros((5,5))

那我想同时把(1,2),(3,4),(0,3)位置的元素分别设置为1,2,3，但是下面不行，

# I expect this to be the same as a[(1,2)] = 1, a[(3,4)] = 2, a[(0,3)] = 3
a[[(1,2),(3,4),(0,3)]] = [1,2,3]

它会提示“数组的索引太多”。我不想涉及效率问题的 for 循环(真正的问题是一个大矩阵，我需要使用一长串随机生成的索引设置或获取该矩阵中的元素)。

Answer 1

您可以执行以下操作:

import numpy as np

rows, cols = zip(*[(1, 2), (3, 4), (0, 3)])

a = np.zeros((5,5))
a[rows, cols] = [1, 2, 3]

print(a)

输出

[[0. 0. 0. 3. 0.]
 [0. 0. 1. 0. 0.]
 [0. 0. 0. 0. 0.]
 [0. 0. 0. 0. 2.]
 [0. 0. 0. 0. 0.]]