📂ジュリア

JuliaでNearstNeighbors.jlを使用して距離を素早く計算する方法

概要

$n$個の座標同士の距離を計算するにあたり、行列を作る必要はなく、単に距離を計算する場合、多次元検索に有利なデータ構造であるk-dツリー¹を使用して速度を上げることができます。NearestNeighbors.jlに関連するアルゴリズムがすべて実装されているので、公式GitHubページを参照してください。

速度比較

pairwise()関数で距離行列の計算に最適化された技術と比較してみましょう。

using Distances
using StatsBase
using Random
using NearestNeighbors

Random.seed!(0);
ε = 0.01
N = 10^4
coordinate = rand(2, N);
state = sample(['S', 'I'], Weights([0.1, 0.9]), N);
S = coordinate[:, state .== 'S']
I = coordinate[:, state .== 'I']

@time sum(pairwise(Euclidean(),S,I) .< ε, dims = 1)
@time kdtree = KDTree(S); contact = length.(inrange(kdtree, I, ε, true))

上記のコードを実行した結果は以下の通りです。一番下の二つのコマンドラインは同じ作業を行いますが、速度差は約500倍になります。実際、k-dツリーで検索する際の時間複雑度は$\log n$で非常に効率的です。

julia> @time sum(pairwise(Euclidean(),S,I) .< ε, dims = 1)
  0.098394 seconds (14 allocations: 69.639 MiB, 8.23% gc Time)
1×9004 Array{Int64,2}:
 0  0  1  0  0  0  0  …  1  0  0  0  0  0

julia> @time kdtree = KDTree(S); contact = length.(inrange(kdtree, I, ε, true))
  0.000213 seconds (22 allocations: 51.609 KiB)
9004-element Array{Int64,1}:
 0
 0
 1
 0
 ⋮
 0
 0
 0

単純な速度の問題を置いておいても、k-dツリーを使用した方法では1次元の配列を返すため、結果物を使用する側面でもさらに便利です。

環境

• OS: Windows
• julia: v1.6.2