Python 如何获取2D numpy数组中所有行的对角线？

给定这个numpy数组

[[200.       202.08165 ]
 [189.60295  190.32434 ]
 [189.19751  188.7867  ]
 [162.15639  164.05934 ]]

我想得到这个数组

[[200.       190.32434 ]
 [189.60295  188.7867  ]
 [189.19751  164.05934 ]]

[[200.       202.08165    187.8392 ]
 [189.60295  190.32434    167.93082]
 [189.19751  188.7867     199.2839 ]
 [162.15639  164.05934    200.92   ]]

[[200.       190.32434    199.2839 ]
 [189.60295  188.7867     200.92   ]]

给定此数组，3列的情况相同

[[200.       190.32434 ]
 [189.60295  188.7867  ]
 [189.19751  164.05934 ]]

[[200.       202.08165    187.8392 ]
 [189.60295  190.32434    167.93082]
 [189.19751  188.7867     199.2839 ]
 [162.15639  164.05934    200.92   ]]

[[200.       190.32434    199.2839 ]
 [189.60295  188.7867     200.92   ]]

我想得到这个数组

[[200.       190.32434 ]
 [189.60295  188.7867  ]
 [189.19751  164.05934 ]]

[[200.       202.08165    187.8392 ]
 [189.60295  190.32434    167.93082]
 [189.19751  188.7867     199.2839 ]
 [162.15639  164.05934    200.92   ]]

[[200.       190.32434    199.2839 ]
 [189.60295  188.7867     200.92   ]]

---

对于任意数量的列和行，有没有矢量化的方法来实现这一点

np.diag

和

np.diagonal

似乎只给了我一条对角线，但我需要把它们全部叠起来。

这似乎是保留对角线元素的一个特殊情况。下面是一个使用

掩蔽的矢量化解决方案

-

def keep_diag(a):    
    m,n = a.shape
    i,j = np.ogrid[:m,:n]
    mask = (i>=j) & ((i-m+n)<=j)
    return a.T[mask.T].reshape(n,-1).T

---

您仍然可以使用

np.diagonal（）

：

---

这回答了你的问题吗@b-fg不幸的是，在这些答案中没有矢量化的方法，我的情况也不同。我认为从我所举的例子中应该可以清楚地看到这一点。