Python:numpyの要素抽出について

Question

前提・実現したいこと

Pythonを使用しています。

• 実現したい内容 …

「　waの最初の行と最後の行だけを抽出したいです。」

実行したコードとその結果

実行したコード

このようなコードを実行しました。

In[]:
wa = np.eye(4, 4, 1) + np.diag((3, 4, 5, 6))
wa[1:-1]

実行した結果

上記のコードを実行した結果、
以下のようになりました。

Out[]:
array([[0., 4., 1., 0.],
       [0., 0., 5., 1.]])

本来得たい結果

望んでいた実行結果は以下の通りです。

Out[]:
array([[3., 1., 0., 0.],
       [0., 0., 0., 6.]])

Answer 1

抽出したい行のインデックスを配列で指定してください。

array[[インデックス1, インデックス2, ...]]

import numpy as np

a = np.eye(4, 4, 1) + np.diag((3, 4, 5, 6))
print(a)
# [[3. 1. 0. 0.]
#  [0. 4. 1. 0.]
#  [0. 0. 5. 1.]
#  [0. 0. 0. 6.]]

b = a[[0, -1]]
print(b)
# [[3. 1. 0. 0.]
#  [0. 0. 0. 6.]]