numpy 3d arrays中的for循環

Question

我想采用這樣的基本 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)

我對上述 output 的問題是將每次迭代放入一個垂直數組，然后移至下一個。 如何使上述代碼 output 類似於我的第二個代碼塊？

抱歉格式不佳，stackoverflow 新手，剛開始學習代碼/numpy。

謝謝！

Answer 1

您可以使用列表理解：

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

Answer 2

您可以使用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]])]

Answer 3

我想這就是你想要做的

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]

Answer 4

另一種方法（大概在您的樓梯中， b在其最后兩個維度中必須具有方形）：

c = b[:,np.arange(b.shape[1]),np.arange(b.shape[2])]
b2 = c.T.reshape(c.shape+(1,))

output：

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

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

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

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