Skip to main content
 首页 » 编程设计

python之使用 numpy 数组将值分配给另一个数组

2025年05月04日25zhwl

我有以下 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] 用于设置元素。