【python】numpyでバイナリサーチをするsearchsorted

　numpy.searchsortedを使うとnumpyでソート済み配列に対するバイナリサーチ（二分探索）を行えます。

　ただし、1次元配列しか受け付けません。まあ、いいか。

　次のように使えます。

>>> import numpy as np
>>> a = np.linspace(0, 100, 1000)
>>> a.shape
(1000,)
>>> np.searchsorted(a, 64)
640
>>> a[640]
64.06406406406407

　つまり、ソート済みの一次元配列を第一引数に、探したい値を第二引数に渡すと、順序を狂わせないで第二引数を挿入できる場所のインデックスを返します。そういう意味では、普通の二分探索とは少し異なります。もしぴったり一致するものがほしいのであれば、返り値のインデックスを得てから自分でチェックしたほうが良いでしょう。

　optionalな引数としてはsideとsorterがあります。

　sideは、"right"と"left"という文字列を指定できます。デフォルトは"left"です。お察しの方もおられると思いますが、配列内のある要素とぴったり一致する値を指定したとき、左右どちらが等号なしの不等号になるかを示します。

side returned index i satisfies
left a[i-1] < v <= a[i]
right a[i-1] <= v < a[i]

>>> np.searchsorted(a, 64.06406406406407, side="left")
640
>>> np.searchsorted(a, 64.06406406406407, side="right")
641

　もう一つの重要な引数はsorterです。名前からはソートを行うcallableでも渡すのかな？　という気がしますが、実際に渡さないといけないのはargsortした結果です。

>>> a = np.random.rand(1000)
>>> idx = a.argsort()
>>> np.searchsorted(a, 0.5, sorter=idx)
539
>>> a[idx][539]
0.5011315014064937

　しょうじき使いづらいと思った。

　他に注意しておくべきことは、渡した配列の範囲内に値が入れられなかった場合の返り値についてです。

>>> a = np.random.rand(1000) # aは上でrandで生成したので[0, 1)の区間に収まる
>>> np.searchsorted(a, -1, sorter=idx)
0  # 0になる
>>> np.searchsorted(a, 2, sorter=idx)
1000  # 配列の要素数n+1になる

　範囲内に収まる保証がないときは、何らかのチェックが必要になるでしょう。

　最後に、一応素直に線形探索する場合と比べた速度を見ておきます。

>>> a = np.arange(10**4)
>>> np.nonzero(a == 5000)[0][0]
5000
>>> np.searchsorted(a, 5000)
5000
>>> import timeit
>>> timeit.timeit(lambda : np.nonzero(a == 5000)[0][0], number=10**5)
1.3569234880014847
>>> timeit.timeit(lambda : np.searchsorted(a, 5000), number=10**5)
0.30346718700093334

　少なくとも数倍の高速化が見込めるでしょう。もちろん計算量が変わってくるので、実際のところどれだけ速くなるかはケースバイケースです。

静かなる名辞

pythonとプログラミングのこと

【python】numpyでバイナリサーチをするsearchsorted