Julia：计算两组点之间最小距离的快速方法

Question

我在 Matrix 中有 5000 个 3D 点A，在矩阵中有另外 5000 个 3D 点B。

对于 中的每个点，A我想找到到 中的点的最小距离B。这些距离应存储在一个包含 5000 个条目的数组中。

到目前为止，我已经有了这个解决方案，运行在大约0.145342 seconds (23 allocations: 191.079 MiB). 我怎样才能进一步改进这一点？

using Distances

A = rand(5000, 3)
B = rand(5000, 3)

mis = @time minimum(Distances.pairwise(SqEuclidean(), A, B, dims=1), dims=2)

Answer 1

这是一种标准方法，因为它将具有更好的时间复杂度（特别是对于较大的数据）：

using NearestNeighbors
nn(KDTree(B'; leafsize = 10), A')[2] .^ 2

两条评论：

• 默认情况下计算欧几里得距离（所以我将其平方）
• 默认情况下，NearestNeigbors.jl 假设观测结果存储在列中（所以我需要B'和A'在解决方案中；如果你的原始数据被转置，则不需要；如此设计的原因是 Julia 使用列主矩阵存储）

Answer 2

使用生成大距离矩阵Distances.pairwise(SqEuclidean(), A, B, dims=1)效率不高，因为与 CPU 缓存和现代 CPU 的计算能力相比，现在的主内存相当慢，并且不会很快变得更好（请参阅“内存墙”）。使用两个基本的嵌套 for 循环来即时计算最小值会更快。此外，可以使用多个核心，使用多个线程更快地计算这一点。

function computeMinDist(A, B)
    n, m = size(A, 1), size(B, 1)
    result = zeros(n)
    Threads.@threads for i = 1:n
        minSqDist = Inf
        @inbounds for j = 1:m
            dx = A[i,1] - B[j,1]
            dy = A[i,2] - B[j,2]
            dz = A[i,3] - B[j,3]
            sqDist = dx*dx + dy*dy + dz*dz
            if sqDist < minSqDist
                minSqDist = sqDist
            end
        end
        result[i] = minSqDist
    end
    return result
end

mis = @time computeMinDist(A, B)

请注意，Julia 解释器默认使用 1 个线程，但可以使用环境变量JULIA_NUM_THREADS=auto或仅使用 flag 运行它来调整--threads=auto。有关更多信息，请参阅多线程文档。

绩效结果

以下是我的 6 核 i5-9600KF 机器（带有两个 5000x3 矩阵）的性能结果：

function computeMinDist(A, B)
    n, m = size(A, 1), size(B, 1)
    result = zeros(n)
    Threads.@threads for i = 1:n
        minSqDist = Inf
        @inbounds for j = 1:m
            dx = A[i,1] - B[j,1]
            dy = A[i,2] - B[j,2]
            dz = A[i,3] - B[j,3]
            sqDist = dx*dx + dy*dy + dz*dz
            if sqDist < minSqDist
                minSqDist = sqDist
            end
        end
        result[i] = minSqDist
    end
    return result
end

mis = @time computeMinDist(A, B)

因此，该实施速度快了 21 倍。结果与少数 ULP 相同。

请注意，代码当然可以使用循环平铺进一步优化，并且可能通过转置 A，B因此 JIT 可以使用SIMD 指令生成更有效的实现。