numpy.ndarrayの一部の値を取得する

numpy.ndarrayの一部の値を取得したいときがある。 例えば、図1の赤枠で囲まれた範囲の値を取得することを考える。

図1：問題の例

このとき、pythonには便利な関数numpy.ix_が用意されている。

numpy.org

これを使うと簡単に所望の値が取得できる。実際にやってみよう。

import numpy as np

A = np.array([[1, 2, 3],
              [4, 5, 6],
              [7, 8, 9]
             ])

A[np.ix_([0, 1], [1, 2])]

"""
出力結果
array([[2, 3],
       [5, 6]])
"""

np.ix_を使うときに気をつけなければいけないことは、np.ix_の引数は全てリストにしなければならないことだ。 リストではなく、int型を与えてみよう。

A[np.ix_(0, [1, 2])]

"""
エラーを吐く
ValueError: Cross index must be 1 dimensional
"""

私はこれにハマって時間を溶かしてしまった。皆さんも気を付けてください。