What does partialsortperm! do?

Yes

So your code would look like

function example(n, k)
    dist = rand(n, n)
    topk = Matrix{Int}(undef, k, n)
    perm = Vector{Int}(undef, n)
    for i = 1 : n
        topk[:, i] .= partialsortperm!(perm, view(dist, :, i), 1:k)
    end
    dist, topk
end

That allocates O((k+1)n) whereas non-mutating partialsortperm version would allocate O(n^2)

1 Like