📂줄리아

줄리아에서 NearstNeighbors.jl로 빠르게 거리 계산하는 법

개요

$n$ 개의 좌표끼리 거리를 계산하는데, 행렬을 만들 필요까지는 없고 단순히 거리만 계산하면 되는 경우 다차원 탐색에 유리한 자료구조인 k-d 트리¹를 사용해 속도를 높일 수 있다. NearestNeighbors.jl에 관련 알고리즘이 모두 구현되어 있으니 공식 깃허브 페이지를 참고하도록 하자.

속도 비교

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