[ノウハウ] numpy arrayで条件に合う要素のみ抜き出す

numpyのndarrayで条件に合う要素のみ抜き出す方法。

ndarrayはサイズ固定なので、行列同士で写像すると１次元化されてしまう。

サンプルコード

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]

 行ごとに比較して論理ベクトルを取得し、元の行に対して射影したものを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]]