通过列表理解设置 numpy 数组值

Question

我可以更改 numpy 矩阵的一组值，以列表形式传递我想更改的索引，例如 matrix[[some first indices], [some second indices]] = 1我正在尝试这样做，但只有一个索引列表，然后将该列表的每个元素作为第一个索引，将整个列表作为另一个索引。这里的例子：import numpy as np#Matrixmatrix = np.zeros((5,5))# Indiceselements = [1,2,3]# Way that worksfor i in elements:    matrix[i, elements] = 1我想将其作为列表理解来执行，但我无法弄清楚如何进行，而且，我不确定这是否是一种好的做法。就像是matrix[[(i,elements) for i in elements]] = 1也在考虑用 itertools 来做，例如matrix[(itertools.permutation(elements, 2))] = 1。但是这些方法中的任何一种都有效，我不确定为什么。

Answer 1

您需要使用行和列数组作为索引访问二维矩阵：

import numpy as np

matrix = np.zeros((5,5))

elements = [1,2,3]

matrix[np.c_[elements], elements] = 1

# array([[0., 0., 0., 0., 0.],

#   [0., 1., 1., 1., 0.],

#   [0., 1., 1., 1., 0.],

#   [0., 1., 1., 1., 0.],

#   [0., 0., 0., 0., 0.]])

Answer 2

您可以创建一个slice从第一个到最后一个元素 (+1) 的elements（前提elements是已排序）：

s = slice(elements[0], elements[-1]+1)

matrix[s, elements]= 1 

# [[0. 0. 0. 0. 0.]

# [0. 1. 1. 1. 0.]

# [0. 1. 1. 1. 0.]

# [0. 1. 1. 1. 0.]

# [0. 0. 0. 0. 0.]]