Skip to content

Instantly share code, notes, and snippets.

@carlobaldassi
Last active August 29, 2015 14:18
Show Gist options
  • Save carlobaldassi/257b34bd99096d35b161 to your computer and use it in GitHub Desktop.
Save carlobaldassi/257b34bd99096d35b161 to your computer and use it in GitHub Desktop.
partial select
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))
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
@sglyon
Copy link

sglyon commented Apr 9, 2015

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

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment