【python】numpyで複数行・列の抜き出し

はじめに

　numpyで複数の行や列を抜き出して使いたいときがある。たとえば機械学習で特徴ベクトルのデータがあって、変数重要度の高い次元だけ取りたい・・・みたいなシチュエーション。

　これには実はやり方があって、簡単にできる。ただし、これを知らないととても不幸な目に遭う（for文頑張って回すしかない、遅い・・・）ので、覚えておくべき。

方法

　次のようにする。

>>> import numpy as np
>>> a = np.array([[1,2,3],[4,5,6],[7,8,9]])
>>> a
array([[1, 2, 3],
       [4, 5, 6],
       [7, 8, 9]])
>>> a[[0,2]] # 行抽出の場合。0行目と2行目だけ抽出
array([[1, 2, 3],
       [7, 8, 9]])
>>> a[:,[0,2]] # 列抽出の場合。0列目と2列目だけ抽出
array([[1, 3],
       [4, 6],
       [7, 9]])

　要するにndarrayのインデックスにリストを指定できるということ。念のため言っておくと、pythonのデフォルトのリストにはこんな機能はない。あくまでもnumpyの独自機能である。

>>> a.tolist()[[0,2]]
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
TypeError: list indices must be integers or slices, not list