web-dev-qa-db-ja.com

NumPy ndarrayから行を選択する

2番目の列の値に基づいて NumPy 配列から特定の行のみを選択したい。たとえば、このテスト配列の2番目の列には1〜10の整数があります。

>>> test = numpy.array([numpy.arange(100), numpy.random.randint(1, 11, 100)]).transpose()
>>> test[:10, :]
array([[ 0,  6],
       [ 1,  7],
       [ 2, 10],
       [ 3,  4],
       [ 4,  1],
       [ 5, 10],
       [ 6,  6],
       [ 7,  4],
       [ 8,  6],
       [ 9,  7]])

2番目の値が4である行のみが必要な場合、簡単です。

>>> test[test[:, 1] == 4]
array([[ 3,  4],
       [ 7,  4],
       [16,  4],
       ...
       [81,  4],
       [83,  4],
       [88,  4]])

しかし、必要な値が複数ある場合に同じ結果を得るにはどうすればよいですか？

必要なリストの長さは任意です。たとえば、2番目の列が2、4、または6であるすべての行が必要な場合があります。

>>> wanted = [2, 4, 6]

私が思いついた唯一の方法は、リスト内包表記を使用してから、これを配列に変換し直すことです。

>>> test[numpy.array([test[x, 1] in wanted for x in range(len(test))])]
array([[ 0,  6],
       [ 3,  4],
       [ 6,  6],
       ...
       [90,  2],
       [91,  6],
       [92,  2]])

NumPy自体でこれを行うより良い方法はありますか？

pythonnumpy
24
Raja Selvaraj
test[numpy.logical_or.reduce([test[:,1] == x for x in wanted])]

NumPyはPythonではなく内部ループを実行するため、結果は元のバージョンよりも高速になります。

16
Amnon

wantedが大きくなると、次のソリューションはAmnonのソリューションよりも高速になります。

_# Much faster look up than with lists, for larger lists:
wanted_set = set(wanted)

@numpy.vectorize
def selected(elmt): return elmt in wanted_set
# Or: selected = numpy.vectorize(wanted_set.__contains__)

print test[selected(test[:, 1])]
_

実際、test配列のみを検索するという利点がありますonce（Amnonの答えのようにlen(wanted)回だけではなく）。また、setsで検索されるPythonの組み込み高速要素を使用します。これは、リストよりもはるかに高速です。また、Numpyの高速ループを使用するため高速です。 in演算子の最適化も得られます。wanted要素が一致したら、残りの要素をテストする必要はありません（Amnonの「論理的または」アプローチとは異なり、 wantedの要素は何であれテストされます）。

別の方法として、次のワンライナーを使用することもできます。これは、配列を1回だけ通過します。

_test[numpy.apply_along_axis(lambda x: x[1] in wanted, 1, test)]
_

ただし、これは各反復で2番目の列の要素を抽出するため、はるかに遅くなります（この回答の最初のソリューションのように1回のパスで要素を抽出するのではなく）。

31
Eric O Lebigot

numpy.in1d はあなたが探しているものです：

print test[numpy.in1d(test[:,1], wanted)]

必要な場合は、簡単に最速のソリューションにする必要があります。さらに、それは最も読みやすいものです、とidが言います。

10
Eelco Hoogendoorn

これは、len（test）= 1000のAmnonのバリアントより2倍高速です。

wanted = (2,4,6)
wanted2 = numpy.expand_dims(wanted, 1)
print test[numpy.any(test[:, 1] == wanted2, 0), :]
0
Antony Hatchkins