如何在没有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.]]