Juliaのsearchsorted関数について

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

searchsorted関数

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

Juliaにはsearchsortedの他にsearchsortedfirstとsearchdortedlastという２つの関数があります。

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

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

aはAbstractVector型でサーチ対象の一次元Array、xはReal型でソートするための実数（Pythonとは異なり一次元配列は引数にできません）。キーワード引数はby、lt、revとなっています。

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

searchsortedfirst関数

searchsortedfirstは

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

を満たす最小の i を返します。また、xがaの中の全要素より大きい時は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 を返します。また、xがaの中の全要素より小さい時は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

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

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

# １番目の要素を見る
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さんのコメントを参照）

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