Juliaのsearchsorted関数について


Juliaでsearchsorted関数を使う機会があったのですが、Juliaの公式documentの説明が結構分かりにくかったので色々調べてみました。

searchsorted関数

Pythonにも同じ機能のライブラリや関数、例えばbisectライブラリ(nagataaaasさんに教えてもらいました)やsearchsorted関数(numpy)はあるのですが、Juliaでは引数が異なってたりするので注意が必要です。Pythonではsearchsorted関数は一次元arrayと実数(もしくはそれから成る一次元array)を引数にとります。この関数はざっくりいうと、その実数に最も近いarrayの要素のインデックスを返すというものです。

Juliaにはsearchsortedの他にsearchsortedfirstsearchdortedlastという2つの関数があります。

基本的に引数は 次のようになっています。

searchsorted(a, x; by=<transform>, lt=<comparison>, rev=false)

aAbstractVector型でサーチ対象の一次元Array、xReal型でソートするための実数(Pythonとは異なり一次元配列は引数にできません)。キーワード引数はbyltrevとなっています。

基本的に配列aは昇順・降順でソート済みであることを前提とします。ソートしてない場合はsortもしくはsort!を使いましょう。

searchsortedfirst関数

searchsortedfirst

a[i-1] < x \le a[i]

を満たす最小の i を返します。また、xaの中の全要素より大きい時はlength(a)+1を返します。

julia> searchsortedfirst([1,3,5,7], 4)
3

julia> searchsortedfirst([1,3,5,7], 9)
5

配列aが降順の時はキーワード引数のrevをtrueにします。このとき

a[i-1] > x \ge a[i]

となる最小のiを返します。

julia> searchsortedfirst([7,5,3,1], 4, rev=true)
3

となります。

searchsortedlast関数

searchsortedlast

a[i] \le x < a[i+1]

を満たす最小の i を返します。また、xaの中の全要素より小さい時は0を返します。

julia> searchsortedlast([1,3,5,7], 4)
2

julia> searchsortedlast([1,3,5,7], -1)
0

同様に降順のときはキーワード引数のrevをtrueにします。このとき

a[i] \ge x > a[i+1]

となる最大のiを返します。

julia> searchsortedlast([7,5,3,1], 4, rev=true)
0

となります。

searchsorted関数

さて最初に戻ってsearchsorted関数を見てみましょう。

julia> searchsorted([1,3,5,7],4)
3:2

julia> typeof(searchsorted([1,3,5,7],4))
UnitRange{Int64}

searchsortedの返り値の型はUnitRange{Int64}になっています(マシンによってはUnitRange{Int32}になるかもしれません)。UnitRangeはよくfor文などで使用される1:10などの型です。
ちなみにUnitRange

struct UnitRange{T<:Real} <: AbstractUnitRange{T}
     start::T
     stop::T
     UnitRange{T}(start, stop) where {T<:Real} = new(start,unitrange_last(start,stop))
end

で型が作られているので、先ほどの3:2のそれぞれの数字は次のように取得できます。

julia> searchsorted([1,3,5,7],4).start
3

julia> searchsorted([1,3,5,7],4).stop
2

ここで気づいた人がいるかもしれませんが、それぞれsearchsortedfirstsearchsortedlastに対応しています。

julia> searchsorted([1,3,5,7],4).start == searchsortedfirst([1,3,5,7],4)
true

julia> searchsorted([1,3,5,7],4).stop == searchsortedlast([1,3,5,7],4)
true

またさらにそれぞれはnumpyのsearchsorted(a,x,side='right')searchsorted(a,x,side='left')にも対応しています(もちろんPythonの場合インデックスが0から始まるという違いがありますが)。

キーワード引数by

キーワード引数のbyは配列aに対し要素間の比較を行う前に施す関数を指定します。

#絶対値を見る
julia> searchsortedfirst([1, -2, 4, -10], -3, by=abs)
3

#1番目の要素を見る
julia> searchsortedfirst([(1,"a"), (3,"b"), (5,"g")], (4,"h"), by=x->x[1])
3

最初に述べた通り配列aはソート済みであることを仮定していました。これは公式ドキュメントにも書いてある通りです。

Return the range of indices of a which compare as equal to x (using binary search) according to the order specified by the by, lt and rev keywords, assuming that a is already sorted in that order.

しかしaがソートされていなくてもsearchsorted関数は使うことはできます。

julia> searchsorted([3,1,6,5], 2)
3:2

julia> searchsortedfirst([3,1,6,5], 2)
3

julia> searchsortedlast([3,1,6,5], 2)
2

ですが先の説明に従うと、searchsortedfirstでは1を返すはずでは?

とりあえずソートしてない配列を渡すのはやめたほうが良さそうな気がします(そもそもソートしてない配列に対してsearchsortedを使いたい場面ってあるのかな?)

antimon2さんより、searchsorted関数binary searchを使っているのでこういうことが起こるそう(確かにdocumentに書いてありますね)。なのでかならずソート済みの配列を用いるべきだし、ソートしてない配列にこの関数を使いたい場面などないみたいです(詳しくはantimon2さんのコメントを参照)

(何か間違い等あればコメントお願いします。)