Python 二次元配列から、任意のデータを探す

Question

np.arrayの二次元配列(要素数は5)から要素①,要素②,要素③,要素④の組み合わせが例えば、1,2,3,4の列の要素⑤を取得する場合はどうすればいいですか？

下記のようにすれば要素⑤の取得は可能ですが、何万回と繰り返す部分なのでもっと速く取得できる方法を教えてください。

for i in array:
    if i[0] == 1 and i[1] == 2 and i[2] == 3 and i[3] == 4:
        print(i[4])

Answer 1

こういうのはどうでしょうか。

import timeit
import numpy as np

A = np.arange(50000).reshape((10000, 5))

def f0(a,b,c,d):
    """質問文の方式
    """
    lst = []
    for i in A:
        if i[0] == 1 and i[1] == 2 and i[2] == 3 and i[3] == 4:
            lst.append(i[4])
    # print(lst)
    lst

def f1(a,b,c,d):
    """私が考えた方式。一列ずつ見て、削っていきます
    こうすると比較の回数がケチれて（本質的には短絡評価と同じことをしています）速いはずです
    """
    tmp = A
    for i, eqv in enumerate([a,b,c,d]):
        tmp = tmp[tmp[:,i] == eqv]
    
    # print(tmp[:,-1])
    tmp[:,-1]  # 計測時にprintしたくなかったので

def f2(a,b,c,d):
    """mkgreiさんの方式
    """
    idx = (A[:, :4]==np.array([a,b,c,d])).all(axis=-1)
    # print(A[idx, -1])
    A[idx, -1]  # 計測時にprintしたくなかったので

print(timeit.timeit(lambda : f0(0,1,2,3), number=1000))
print(timeit.timeit(lambda : f1(0,1,2,3), number=1000))
print(timeit.timeit(lambda : f2(0,1,2,3), number=1000))
"""
私の環境だと、
3.3318773460000557
0.06133379499988223
0.27946032499994544
くらい
"""

Comments

[mkgrei]

なるほど、コケを落としていくように絞り込んでいくんですね。
4倍以上も速い…

[hayataka2049]

これだと相手のデータがでかいほど効きます。小さいとどこかでオーバーヘッドに負けて逆転するはずです。
本当はぜんぶnumpy上でやろうと思ったんですが、うまく書けなかったのでとりあえずただのforで

[O.K.]

回答ありがとうございます。
tmp = tmp[tmp[:,i] == eqv]
の部分がよくわかりません。
どうなっているのですか？

[hayataka2049]

enumerateの機能を確認するとわかりますが、そのforループは
tmp = tmp[tmp[:,0] == a]
tmp = tmp[tmp[:,1] == b]
tmp = tmp[tmp[:,2] == c]
tmp = tmp[tmp[:,3] == d]
と4行書いたのと等価です。

[O.K.]

それはわかっていました。
tmpという変数に代入しているtmp[tmp[:,0] == a]は何ですか？
[ ]の中に等式があるのを見たことがないのですが...

[hayataka2049]

numpy配列とスカラーの比較をすると、booleanの配列になります。
booleanの配列をインデックスに入れると、Trueのところだけ取り出せます。
https://docs.scipy.org/doc/numpy-1.13.0/user/basics.indexing.html#boolean-or-mask-index-arrays

[mkgrei]

よく見たらこの速すぎるテストケースは……

まあランダムでも２倍以上速いんですけどね。

[hayataka2049]

あ、確かにこれだと1万回+1回+1回+1回の比較ですから、ちょっとずるかったですね・・・
現実的な速度は毎回どれだけ削れていくか次第です

Answer 2

import numpy as np 

a = np.random.randint(2, size=(50, 3))

idx = (a[:, :2]==np.array([0,1])).all(axis=-1)

print(a[idx, -1])

Comments

[O.K.]

回答ありがとうございます。
二行目以降が何をしているのかわからないのですが、解説していただけませんか？

[mkgrei]

0か1の50×3の配列を先に作っています。
これはご質問されているものの縮小版です。
1番目と2番目の要素が特定の数字の場合、3番目の要素を得る、というように読み替えてください。

次にa[:, :2]==np.array([0,1])のところで、aの50×2の部分を取り出して[0,1]と比較します。
それぞれの行において、1番目と2番目が0と1ならTrueになります。
例えば[0,0]なら[True,False]、[1,1]なら[False,True]などになります。
挙動が不明瞭ならprintで書き出してみてください。

これに.all(axis=-1)をすると、それぞれの行で全ての要素がTrueの場合に限りTrueになります。
例えば、[[True,True],[True,False]]なら[True,False]になります。
ここで次元が1つ落ちます。

最後に、numpy.arrayの機能で、配列と同じ長さのbool配列を入れると、Trueの部分だけ取り出すことができることを利用します。
ですので、a[idx, -1]というのはidxがTrueの行の1番最後の要素、今の場合3番目の要素を取り出すということになります。

これらを応用することでご質問されているような要素を高速に抜き出すことができます。

[hayataka2049]

対案投げてみました

Answer 3

何も考えないで自然に書いたら他の人より100倍以上遅いんですけどなんなんですかね

def f3(a, b, c, d):
    result = []
    c = np.array([1, 2, 3, 4])
    for r in A:
        if np.array_equal(r[0:4], c):
            result.append(r[4])

指摘お待ちしてまーす

今見たらキモかったから書き直した

import timeit
import numpy as np

A = np.random.randint(low=1, high=5, size=(10000, 5))


def fao_fun(A):
    result = []
    t = np.array([1, 2, 3, 4])
    for row in A:
        if np.array_equal(row[0:4], t):
            result.append(row[4])

print(timeit.timeit(lambda: fao_fun(A), number=1))

0.10979769402183592

やっぱり他の人の100倍位遅いね
なんでかな

Answer 4

最近覚えたmemoryviewで。

import numpy as np
view = memoryview(np.array([95, 96, 97, 98]))
view_size = len(view)
A = np.arange(50000).reshape((10000, 5))

for row in A:
    if memoryview(row)[:view_size] == view:
        print(row[-1])

追記

ans = tuple(v[-1] for v in map(memoryview, A) if v[:view_size] == view)
print(ans)

追記２：恥ずかしいくらい遅いので本気を出す

慣れないnumpyを使わずに2dのtuple->listにして素直に書いた方が結果自己記録的に速かった。

import numpy as np

t = [95, 96, 97, 98]
A = tuple(map(list, np.arange(50000).reshape((10000, 5))))
tuple(last for *first, last in A if first == t)

Comments

[YouheiSakurai]

あぁ、それにしてもnumpy速い。