我有以下 numpy 数组 matrix ,
matrix = np.zeros((3,5), dtype = int)
array([[0, 0, 0, 0, 0],
[0, 0, 0, 0, 0],
[0, 0, 0, 0, 0]])
假设我也有这个 numpy 数组 indices
indices = np.array([[1,3], [2,4], [0,4]])
array([[1, 3],
[2, 4],
[0, 4]])
问题如何将 1 分配给 matrix 中的元素,其中它们的索引由 indices 数组。预计会进行矢量化实现。
为了更清晰,输出应该如下所示:
array([[0, 1, 0, 1, 0], #[1,3] elements are changed
[0, 0, 1, 0, 1], #[2,4] elements are changed
[1, 0, 0, 0, 1]]) #[0,4] elements are changed
请您参考如下方法:
这是使用 NumPy's fancy-indexing 的一种方法-
matrix[np.arange(matrix.shape[0])[:,None],indices] = 1
解释
我们使用 np.arange(matrix.shape[0]) 创建行索引 -
In [16]: idx = np.arange(matrix.shape[0])
In [17]: idx
Out[17]: array([0, 1, 2])
In [18]: idx.shape
Out[18]: (3,)
列索引已作为 indices -
In [19]: indices
Out[19]:
array([[1, 3],
[2, 4],
[0, 4]])
In [20]: indices.shape
Out[20]: (3, 2)
我们来做一个行列索引的形状示意图,idx和indices -
idx (row) : 3
indices (col) : 3 x 2
为了使用行索引和列索引对输入数组 matrix 进行索引,我们需要使它们可以相互广播。一种方法是在 idx 中引入一个新轴,通过将元素插入第一个轴并允许使用 将单例 dim 作为最后一个轴,使其成为 ,如下图-2D idx[:,None]
idx (row) : 3 x 1
indices (col) : 3 x 2
在内部,idx 将被广播,就像这样 -
In [22]: idx[:,None]
Out[22]:
array([[0],
[1],
[2]])
In [23]: indices
Out[23]:
array([[1, 3],
[2, 4],
[0, 4]])
In [24]: np.repeat(idx[:,None],2,axis=1) # indices has length of 2 along cols
Out[24]:
array([[0, 0], # Internally broadcasting would be like this
[1, 1],
[2, 2]])
因此,来自 idx 的广播元素将用作来自 indices 的行索引和列索引,用于索引到 matrix 中,用于设置元素它。因为,我们有 -
idx = np.arange(matrix.shape[0]),
因此,我们最终会得到 -
matrix[np.arange(matrix.shape[0])[:,None],indices] 用于设置元素。
