Python 使用Numpy快速提取多维数组下三角的方法？

假设形状为4D数组

4,4,1,1

。目标是提取

（n，m）

数组的下三角

目前，以下代码依赖于for循环来完成此任务

import numpy as np
lw_up_pair = np.tril_indices(4, -1)
arr=np.zeros((4,4,1,1))

arr[1,:1,:,0]=1
arr[2,:2,0,0]=2
arr[3,:3,0,0]=3

ext=[arr [x, i, :, 0] for x,i in zip(lw_up_pair [0],lw_up_pair [1])] # expected output

但是，我想知道是否有更快的或

numpy

内置的方法可以实现类似的结果？

只需执行以下操作：

rows, cols = lw_up_pair
ext = arr[rows, cols, :, 0]
print(ext)

输出

[[1.]
 [2.]
 [2.]
 [3.]
 [3.]
 [3.]]

您还可以执行以下操作：

arr[lw_up_pair][:, 0]

array([[1.],
       [2.],
       [2.],
       [3.],
       [3.],
       [3.]])

---

由于我们正在迭代改进索引表达式：

arr[(*np.tril_indices(4, -1), slice(None), 0)]

---

zip是

np.stack（lw\u up\u pair，-1）

聪明。但是，潜在的缺点是您不能分配给这个表达式。您可以始终只分配给

arr[lw\u up\u pair]

，因为其余的只是为了获得正确的维数，并且您可以非常肯定地获得正确的广播。