julia - 有什么方法可以免费找到 K 最近邻分配？

Question

我需要这个来进行模拟研究。

MWE:

x = rand(10,4)
y = rand(5,4)

对于 y 中的每一行，我想在 x 中找到其 5-nn 的索引，即结果应该是一个 5×5 的索引矩阵。

Answer 1

事实证明这是不完整的，但我还是会发布我的尝试。

在没有分配的情况下将矩阵“重新解释”为向量的向量在概念上很简单，但需要实现新的数组类型。 JuliennedArrays.jl 提供了这种类型的 Sliced。

IHMO 最简单的实现是这个:

mapslices(y, dims=2) do row
    partialsortperm(Slices(x, 2), 1:5, by=x -> norm(x - row))
end 

仍然分配一些东西；这必须至少是 partialsortperm 和中间行使用的索引向量。

我试图在这个函数中去掉它:

function knnslice!(result, x, y, k)
    result_sliced = Slices(result, 2)
    x_sliced = Slices(x, 2)
    y_sliced = Slices(y, 2)
    indices = collect(axes(x, 1))
    for i in eachindex(result_sliced, y_sliced)
        result_sliced[i] .= partialsortperm!(indices, x_sliced, 1:k, by=x -> norm(x - y_sliced[i]))
    end
    return result
end
knnslice(x, y, k) = knnslice!(similar(x, Int, size(y, 1), k), x, y, k)

但结果几乎没有改进，至少在与示例数据大小的数组进行比较时是这样。我不确定如何通过这种实现方式进一步降低这种情况。

缺少的部分是一个直接作用于切片的 sortperm 实现。对于较小的 k，这应该可以通过对 x 进行一次迭代并将结果行维护为该大小的缓冲区(甚至是小堆)而不是执行部分排序来实现。像这样的东西:

function knnslice!(result, x, y, k)
    for (i_r, i_y) in zip(axes(result, 1), axes(y, 1))
        result_row = @view(result[i_r, :])
        fill!(result_row, 1)
        f(r) = norm(@view(x[r, :]) - @view(y[i_y, :]))
        for j_x in axes(x, 1)
            heappush!(result_row, j_x; by=f)
        end
    end
    return result
end

heappush! 应该按 by 的顺序插入有界最小堆(类似于 Python 中的 heapq，但保持队列的大小固定).