partial select

partselect_example.jl

using PartSort

# we have some array named `ndata`
# we need the indices of the lowest `k` elements
# we use a preallocated vector `sp` so that we can do this in a loop

nn = length(ndata)
sp = Array(Int, nn)

# this can go in a loop
@inbounds for i = 1:nn; sp[i] = i; end
sort!(sp, PartialQuickSort(k), Base.Perm(Base.Forward,ndata))

PartSort.jl

module PartSort

using Base.Order

import Base: sort!, Sort.Algorithm

export PartialQuickSort

immutable PartialQuickSort <: Algorithm
    k::Int
end

Base.Sort.Float.fpsort!(v::AbstractVector, a::PartialQuickSort, o::Ordering) = sort!(v, 1, length(v), a, o)

function sort!(v::AbstractVector, lo::Int, hi::Int, a::PartialQuickSort, o::Ordering)
    k = a.k
    while lo < hi
        hi-lo <= Base.Sort.SMALL_THRESHOLD && return sort!(v, lo, hi, Base.Sort.SMALL_ALGORITHM, o)
        pivot = v[(lo+hi)>>>1]
        i, j = lo, hi
        while true
            while lt(o, v[i], pivot); i += 1; end
            while lt(o, pivot, v[j]); j -= 1; end
            i <= j || break
            v[i], v[j] = v[j], v[i]
            i += 1; j -= 1
        end
        if lo < j
            if j - lo <= k
                sort!(v, lo, j, QuickSort, o)
            else
                sort!(v, lo, j, PartialQuickSort(k), o)
            end
        end
        jk = min(j, lo + k - 1)
        if (i - lo + 1) <= k
            k -= j - lo + 1
            lo = i
        else
            break
        end
    end
    return v
end

end

This is great. Thanks for sharing. I'll test it out in my algorithm and report a performance comparison.