使用循环修改 np.diag 的对角数组

Question

我一直在尝试查找 np.diag_indices 的工作原理以及它们的示例，但是它的文档有点简单。我知道这会通过您的矩阵创建一个对角线数组，但是我想更改对角线数组（我正在考虑使用循环来更改它的尺寸或沿着这些线的东西）。

I.E. 假设我们有一个 3x2 矩阵：

[[1 2]
 [3 4]
 [5  6]]

现在如果我使用 np.diag_indices 它将形成一个从 (0,0) 开始并经过 (1,1) 的对角线数组。

[1 4]

但是，我希望这个对角线阵列向下移动一个。所以现在它从 (0,1) 开始并经过 (1,2)。

[3 6]

但是，np.diag_indices 只有 2 个参数，从外观上看，这两个参数都不能让我这样做。我是否使用错误的工具来尝试实现这一目标？如果是这样，我可以使用哪些工具来创建穿过我的矩阵的不断变化的对角线数组？ （我正在寻找也适用于 200x50 等较大矩阵的东西）。

Answer 1

diag_indices的代码很简单，简单到没用过：

idx = arange(n)
return (idx,) * ndim

In [68]: np.diag_indices(4,2)                                                   
Out[68]: (array([0, 1, 2, 3]), array([0, 1, 2, 3]))

它只返回一个数组元组，arange 重复了n 次。它对于索引方阵的主对角线很有用，例如

In [69]: arr = np.arange(16).reshape(4,4)                                       
In [70]: arr                                                                    
Out[70]: 
array([[ 0,  1,  2,  3],
       [ 4,  5,  6,  7],
       [ 8,  9, 10, 11],
       [12, 13, 14, 15]])
In [71]: arr[np.diag_indices(4,2)]                                              
Out[71]: array([ 0,  5, 10, 15])

该应用程序使用两个形状匹配的数组直接索引。

它适用于其他形状 - 如果它们足够大。

np.diag 应用于同一个数组做同样的事情：

In [72]: np.diag(arr)                                                           
Out[72]: array([ 0,  5, 10, 15])

但它也允许偏移：

In [73]: np.diag(arr, 1)                                                        
Out[73]: array([ 1,  6, 11])

===

使用diag_indices 进行索引确实允许我们更改对角线：

In [78]: arr[np.diag_indices(4,2)] += 10                                        
In [79]: arr                                                                    
Out[79]: 
array([[10,  1,  2,  3],
       [ 4, 15,  6,  7],
       [ 8,  9, 20, 11],
       [12, 13, 14, 25]])

====

但我们不必使用diag_indices 来生成所需的索引数组：

In [80]: arr = np.arange(1,7).reshape(3,2)                                      
In [81]: arr                                                                    
Out[81]: 
array([[1, 2],
       [3, 4],
       [5, 6]])

从第一 2 行和列中选择值：

In [82]: arr[np.arange(2), np.arange(2)]                                        
Out[82]: array([1, 4])
In [83]: arr[np.arange(2), np.arange(2)] += 10                                  
In [84]: arr                                                                    
Out[84]: 
array([[11,  2],
       [ 3, 14],
       [ 5,  6]])

对于不同的行选择：

In [85]: arr[np.arange(1,3), np.arange(2)] += 20                                
In [86]: arr                                                                    
Out[86]: 
array([[11,  2],
       [23, 14],
       [ 5, 26]])

advanced indexing 的相关文档部分与整数数组：https://numpy.org/doc/stable/reference/arrays.indexing.html#purely-integer-array-indexing