Julia: Given an iteratble `itr`, is there an efficient data structure and algorithm to obtain the index of the top `n` values of `itr`?

Viewed 62

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?

2 Answers

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)

Related