如何在没有for-loop的情况下通过多个索引在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.]]