Consider dist = rand(4_000_000) and I wish to find the index of the top n values. For top 1 it is easy with
argmax(dist)
But if there an function to find the top n?
Consider dist = rand(4_000_000) and I wish to find the index of the top n values. For top 1 it is easy with
argmax(dist)
But if there an function to find the top n?
If you want just the indexes without actually sorting the list use partialsortperm:
julia> partialsortperm(dist,1:5)
5-element view(::Array{Int64,1}, 1:5) with eltype Int64:
1015972
672133
1519815
2724755
428060
Note that if you are doing this many times in a row you can avoid re-allocating the vector of indices by using partialsortperm!.
If you actually need the elements use partialsort, the fastest will be sorting in-place with partialsort!:
julia> partialsort!(dist,1:5)
5-element view(::Array{Float64,1}, 1:5) with eltype Float64:
3.7341092440357215e-8
8.483718505480908e-8
2.1386679072143977e-7
2.2127291532392235e-7
4.3546665384752714e-7
Here is an alternative accepting just an iterable itr that does not require it to be an AbstractArray (I assume itr produces floats for simplicity):
function topn(itr, n)
topvalues = Vector{Tuple{Float64, Int}}() # assuming itr produces floats
@inbounds for (i,v) in enumerate(itr)
if i <= n
push!(topvalues, (v, i))
i == n && sort!(topvalues)
else
if v > first(topvalues[1])
j = searchsortedlast(topvalues, v, by = first) # assuming itr produces floats
for k in 1:j-1
topvalues[k] = topvalues[k+1]
end
topvalues[j] = (v, i)
end
end
end
return sort!(topvalues, rev=true)
end
Now let us check its performance:
julia> using BenchmarkTools
julia> dist = rand(4_000_000);
julia> @btime partialsortperm($dist, 1:5, rev=true)
69.583 ms (4 allocations: 30.52 MiB)
5-element view(::Array{Int64,1}, 1:5) with eltype Int64:
3482486
404839
511215
3082421
2600778
julia> @btime topn($dist, 5)
4.079 ms (5 allocations: 464 bytes)
5-element Array{Tuple{Float64,Int64},1}:
(0.9999999876010581, 3482486)
(0.9999998674122681, 404839)
(0.9999998410867581, 511215)
(0.9999995172666474, 3082421)
(0.9999992939127473, 2600778)
julia> sort!(dist); # make dist ascending
julia> @btime partialsortperm($dist, 1:5, rev=true)
45.634 ms (4 allocations: 30.52 MiB)
5-element view(::Array{Int64,1}, 1:5) with eltype Int64:
4000000
3999999
3999998
3999997
3999996
julia> @btime topn($dist, 5)
46.660 ms (5 allocations: 464 bytes)
5-element Array{Tuple{Float64,Int64},1}:
(0.9999999876010581, 4000000)
(0.9999998674122681, 3999999)
(0.9999998410867581, 3999998)
(0.9999995172666474, 3999997)
(0.9999992939127473, 3999996)
julia> sort!(dist, rev=true); # make dist descending
julia> @btime partialsortperm($dist, 1:5, rev=true)
41.199 ms (4 allocations: 30.52 MiB)
5-element view(::Array{Int64,1}, 1:5) with eltype Int64:
1
2
3
4
5
julia> @btime topn($dist, 5)
4.075 ms (5 allocations: 464 bytes)
5-element Array{Tuple{Float64,Int64},1}:
(0.9999999876010581, 1)
(0.9999998674122681, 2)
(0.9999998410867581, 3)
(0.9999995172666474, 4)
(0.9999992939127473, 5)
So as you can see if n is small this approach is just faster even for AbstractArrays (but the main point is that this method works for any iterable itr).
If n would be large and itr were already sorted in ascending order then partialsortperm with rev=true will be faster (as on-line method then has to update the solution in every iteration which is unavoidable with on-line approach). But I assume that a typical use case is for small n and somewhat shuffled itr.
(if you had to handle the large n case and be more robust you could consider using a heap instead of just a sorted array, but I have not benchmarked it)