JPCB レシピ 2.1 配列中の最小要素のインデックスを取得する

オライリー・ジャパンから刊行されている「Juliaプログラミングクックブック」のレシピを写経してみました．実行したJuliaのバージョンは以下です．

Julia Version 1.5.3

このレシピでは，配列中の最小要素のインデックス返す関数，特に最小要素が複数ある場合に一様ランダムに選択したインデックスを１つ返す関数について議論しています．

まず，配列中のすべての最小要素のインデックスを返す関数は次のように実装されます．

# 最小要素のインデックスをすべて(配列として)返す
function allargmin(a)
    m = minimum(a)
    filter(i -> a[i] == m, eachindex(a))
end

#%% 実行
ar = [1,2,3,1,2,3]
allargmin(ar) # [1, 4]

eachindex()は配列のインデックスをコレクションとして返す関数のようです．

# eachindexは配列のインデックスを返す
for i = eachindex(ar)
    println(i)
end
#= 実行結果
1
2
3
4
5
6
=#

本題にもどって，allargmin()を使って最小要素を一様ランダムに選択して返す関数を実装します．

#%% 最小要素を一様ランダムに選択して返す関数
randargmin1(a) = rand(allargmin(a))

そして，同様の関数を別実装した以下の関数も用意します．

function randargmin2(a)
    indices = eachindex(a)
    y = iterate(indices)
    y === nothing && throw(ArgumentError("collection must be non-empty"))
    (idx, state) = y
    minval = a[idx]
    bestidx = idx
    bestcount = 1
    y = iterate(indices, state)
    while y !== nothing
        (idx, state) = y
        curval = a[idx]
        if isless(curval, minval)
            minval = curval
            bestidx = idx
            bestcount = 1
        elseif isequal(curval, minval)
            bestcount += 1
            rand() * bestcount < 1 && (bestidx = idx)
        end
        y = iterate(indices, state)
    end
    bestidx
end

早速実行してみましょう．

using StatsBase
x = [1, 2, 3, 1, 2, 3, 1, 1]
countmap([randargmin1(x) for i in 1:10^6])
# Dict{Int64,Int64} with 4 entries:
#   7 => 250606
#   4 => 249997
#   8 => 249730
#   1 => 249667
countmap([randargmin2(x) for i in 1:10^6])
# Dict{Int64,Int64} with 4 entries:
#   7 => 249765
#   4 => 249842
#   8 => 250474
#   1 => 249919

乱数を使用しているので結果は一致しませんが，同じ機能が実現できていることが確認できました．
countmap()は要素のヒストグラムをとってくれる関数のようです．便利ですね．

実行時間についても比較してみましょう．

x = rand(1:10, 1000)
using BenchmarkTools
@btime randargmin1($x)
#   1.233 μs (2 allocations: 8.81 KiB)
# 724
@btime randargmin2($x)
#   1.670 μs (0 allocations: 0 bytes)
# 272

マクロでは引数を$xにするのが一般的なのでしょうか？xを引数にしても問題なく実行できましたが，書籍の記述は$xになっています．

興味深いことに，書籍ではrandargmin2($x)の方が17倍も速いのに対し，手元の環境ではrandargmin1($x)の方が速くなっています．書籍はJulia 1.2.0での実行結果のようなので，バージョンが上がって最適化が進んだようですね．実装方法としてrandargmin1($x)の方があきらかにコード量が少ないので，こちらの実行速度が速くなったことは素晴らしいと思います．
またメモリ量に関して，randargmin1($x)は書籍では18.05KB使用しており，手元の環境よりもメモリ使用量が多いです．

その他：isless(), isequal()について．
Juliaでは大小比較や等号比較をする際にはこの関数を使うべきとのことです．演算子による比較をした場合，以下のようにいくつかの値で振る舞いが異なるようです．

0.0 == -0.0, isequal(0.0, -0.0)
# (true, false)
-0.0 < 0.0, isless(-0.0, 0.0)
# (false, true)
NaN == NaN, isequal(NaN, NaN)
# (false, true)
NaN < NaN, isless(NaN, NaN)
# (false, false)

「Juliaプログラミングクックブック」のレシピ2.1を実装してみました．知らない文法やバージョンの違いによる結果の差異など，学ぶことが多かったので，今後も写経を続けていこうと思います．

https://www.oreilly.co.jp/books/9784873118895/