Python 用于numpy 3d阵列中的循环

我想要一个基本的3d阵列，如下所示：

b = np.arange(1,101).reshape(4,5,5)
b

然后我想取第一个索引，然后像楼梯一样向下工作

b1 = [b[0:,0,0],b[0:,1,1],b[0:,2,2],b[0:,3,3],b[0:,4,4]]
b1 = np.asarray(b1)
b1 = np.transpose(b1)
b1

上面的代码看起来不正确，我宁愿使用循环。 这就是我到目前为止所做的：

for i in range(0,5):
    b2 = b[0:,i,i]
    b2 = np.asarray(b2)
    b2 = b2.reshape(4,1)
    print(b2)

我对上述输出的问题是，它将每个迭代放在一个垂直数组中，然后移动到下一个。如何使上面的代码输出类似于我的第二块代码

为糟糕的格式表示歉意，这是stackoverflow的新手，刚刚开始学习代码/numpy

---

谢谢

您可以使用列表理解：

b1 = [b[0:,i,i] for i in range(5)]
b1 = np.asarray(b1)
b1 = np.transpose(b1)

---

您可以使用

einsum

一次提取所有楼梯：

>>> np.einsum("ijj->ij",b) 
array([[  1,   7,  13,  19,  25],
       [ 26,  32,  38,  44,  50],
       [ 51,  57,  63,  69,  75],
       [ 76,  82,  88,  94, 100]])

然后拆分为列：

>>> np.split(np.einsum("ijj->ij",b),np.arange(1,5),1)
[array([[ 1],
       [26],
       [51],
       [76]]), array([[ 7],
       [32],
       [57],
       [82]]), array([[13],
       [38],
       [63],
       [88]]), array([[19],
       [44],
       [69],
       [94]]), array([[ 25],
       [ 50],
       [ 75],
       [100]])]

---

我想这就是你想要做的

import numpy as np

b = np.arange(1,101).reshape(4,5,5)

result = np.diag(b[0])

print(result)

结果:

[ 1  7 13 19 25]

另一种方法是（大概在您的楼梯中，

b

的最后两个维度必须为方形）：

输出：

[[[  1]
  [ 26]
  [ 51]
  [ 76]
  [  7]]

 [[ 32]
  [ 57]
  [ 82]
  [ 13]
  [ 38]]

 [[ 63]
  [ 88]
  [ 19]
  [ 44]
  [ 69]]

 [[ 94]
  [ 25]
  [ 50]
  [ 75]
  [100]]]

[[[  1]
  [ 26]
  [ 51]
  [ 76]
  [  7]]

 [[ 32]
  [ 57]
  [ 82]
  [ 13]
  [ 38]]

 [[ 63]
  [ 88]
  [ 19]
  [ 44]
  [ 69]]

 [[ 94]
  [ 25]
  [ 50]
  [ 75]
  [100]]]