numpyの配列から条件に合う要素のみ抜き出す

pythonで数値を扱っていると必ずお世話になる numpy 。
しかしnumpyの配列であるndarrayから条件に合致するものだけ抽出する場合、少し工夫が必要。
今回はその工夫の方法を紹介する。

python配列特徴

pythonの配列は条件を指定すると、条件に合うものをTrue、合わないものをFalseで返すBoolean配列になる。
さらに元の配列から値を取得する際に、インデックスとして同じサイズのBoolean配列を指定すると、Trueで指定した値だけ含まれた配列を返す。

サンプルコード

import numpy as np
a = np.array([1,2,3,4])
print(a)

af = (a>=2)
print(af)

ad = a[af]
print(ad)

実行結果

[1 2 3 4]
[False  True  True  True]
[2 3 4]

多次元配列を写像すると１次元配列になってしまう

同じ方法を多次元配列に適用すると、１次元配列になって返ってきてしまう。

サンプルコード

a = np.array([
     [1,2,3,4]
     ,[2,3,4,5]
     ,[3,4,5,6]
     ])
a4f = (a>=4)
a4 = a[a4f]
print(a4)

実行結果

[4 4 5 4 5 6]

原因はnumpyの仕様

numpyの配列であるndarrayは配列というより 行列 のようなもので、行毎に列の数を変えるなんて事はできない。
Boolean配列を用いて写像する際は行によって列の数が変わる可能性があるので、結果を１次元配列が返してしまう。

次元数を変えずに写像するには

次元数を変えずに写像するには numpy から list 形式にしてしまうのが手っ取り早い。

サンプルコード

a = np.array([
     [1,2,3,4]
     ,[2,3,4,5]
     ,[3,4,5,6]
     ])

a4l = []
i = 0
for r in a:
    a4l.append(r[r >= 4].tolist())
    i+=1

print(a4l)

実行結果

[[4], [4, 5], [4, 5, 6]]