use ArrayOfArrays for return value to reduce the number of allocated arrays

Project.toml

@@ -3,13 +3,15 @@ uuid = "b8a86587-4115-5ab1-83bc-aa920d37bbce"
 version = "0.4.18"
 
 [deps]
+ArraysOfArrays = "65a8f2f4-9b39-5baf-92e2-a9cc46fdf018"
 Distances = "b4f34e82-e78d-54a5-968a-f98e89d6e8f7"
 StaticArrays = "90137ffa-7385-5640-81b9-e52037218182"
 
 [compat]
 Distances = "0.9, 0.10"
 StaticArrays = "0.9, 0.10, 0.11, 0.12, 1.0"
 julia = "1.6"
+ArraysOfArrays = "0.6"
 
 [extras]
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"

src/NearestNeighbors.jl

@@ -4,6 +4,7 @@ using Distances
 import Distances: PreMetric, Metric, result_type, eval_reduce, eval_end, eval_op, eval_start, evaluate, parameters
 
 using StaticArrays
+using ArraysOfArrays
 import Base.show
 
 export NNTree, BruteTree, KDTree, BallTree, DataFreeTree

src/inrange.jl

@@ -15,10 +15,13 @@ function inrange(tree::NNTree,
     check_input(tree, points)
     check_radius(radius)
 
-    idxs = [Vector{Int}() for _ in 1:length(points)]
+    idxs = VectorOfArrays{Int, 1}()
+    idx = Int[]
 
     for i in 1:length(points)
-        inrange_point!(tree, points[i], radius, sortres, idxs[i])
+        inrange_point!(tree, points[i], radius, sortres, idx)
+        push!(idxs, idx)
+        resize!(idx, 0)
     end
     return idxs
 end
@@ -66,11 +69,14 @@ function inrange_matrix(tree::NNTree{V}, points::AbstractMatrix{T}, radius::Numb
     check_input(tree, points)
     check_radius(radius)
     n_points = size(points, 2)
-    idxs = [Vector{Int}() for _ in 1:n_points]
+    idxs = VectorOfArrays{Int, 1}()
+    idx = Int[]
 
     for i in 1:n_points
         point = SVector{dim,T}(ntuple(j -> points[j, i], Val(dim)))
-        inrange_point!(tree, point, radius, sortres, idxs[i])
+        inrange_point!(tree, point, radius, sortres, idx)
+        push!(idxs, idx)
+        resize!(idx, 0)
     end
     return idxs
 end

src/knn.jl

@@ -18,10 +18,17 @@ function knn(tree::NNTree{V}, points::AbstractVector{T}, k::Int, sortres=false,
     check_input(tree, points)
     check_k(tree, k)
     n_points = length(points)
-    dists = [Vector{get_T(eltype(V))}(undef, k) for _ in 1:n_points]
-    idxs = [Vector{Int}(undef, k) for _ in 1:n_points]
+    dists = VectorOfArrays{get_T(eltype(V)), 1}()
+    idxs = VectorOfArrays{Int, 1}()
+    dist = zeros(get_T(eltype(V)), k)
+    idx = zeros(Int, k)
+
     for i in 1:n_points
-        knn_point!(tree, points[i], sortres, dists[i], idxs[i], skip)
+        knn_point!(tree, points[i], sortres, dist, idx, skip)
+        push!(dists, dist)
+        push!(idxs, idx)
+        fill!(dist, 0)
+        fill!(idx, 0)
     end
     return idxs, dists
 end
@@ -77,12 +84,18 @@ function knn_matrix(tree::NNTree{V}, points::AbstractMatrix{T}, k::Int, ::Val{di
     check_input(tree, points)
     check_k(tree, k)
     n_points = size(points, 2)
-    dists = [Vector{get_T(eltype(V))}(undef, k) for _ in 1:n_points]
-    idxs = [Vector{Int}(undef, k) for _ in 1:n_points]
+    dists = VectorOfArrays{Float64, 1}()
+    idxs = VectorOfArrays{Int, 1}()
+    dist = zeros(Float64, k)
+    idx = zeros(Int, k)
 
     for i in 1:n_points
         point = SVector{dim,T}(ntuple(j -> points[j, i], Val(dim)))
-        knn_point!(tree, point, sortres, dists[i], idxs[i], skip)
+        knn_point!(tree, point, sortres, dist, idx, skip)
+        push!(dists, dist)
+        push!(idxs, idx)
+        fill!(dist, 0)
+        fill!(idx, 0)
     end
     return idxs, dists
 end

test/test_knn.jl

@@ -133,8 +133,8 @@ end
     points = rand(SVector{3, Float64}, 100)
     kdtree = KDTree(points)
     idxs, dists = knn(kdtree, view(points, 1:10), 3)
-    @test idxs isa Vector{Vector{Int}}
-    @test dists isa Vector{Vector{Float64}}
+    @test eltype(idxs) <: AbstractVector{Int}
+    @test eltype(dists) <: AbstractVector{Float64}
 end
 
 @testset "mutating" begin