Mercurial > lbo > hg > juliaplay
changeset 26:015fe10a59cd
arrays: quickselect / kth smallest
author | Lewin Bormann <lbo@spheniscida.de> |
---|---|
date | Sat, 04 Mar 2023 11:31:42 +0100 |
parents | 203bb5bbeea5 |
children | 668adaf0313f |
files | julia/arrays.jl |
diffstat | 1 files changed, 63 insertions(+), 0 deletions(-) [+] |
line wrap: on
line diff
--- a/julia/arrays.jl Fri Mar 03 12:49:11 2023 +0100 +++ b/julia/arrays.jl Sat Mar 04 11:31:42 2023 +0100 @@ -370,3 +370,66 @@ end mat end + +function kth_smallest(a::AbstractVector{T}, k=3) where {T<:Integer} + buf = collect(a[begin:k]) # copy first few elements + sort!(buf) + + fit_in(x) = begin + if x >= buf[end] + return + else + i = findnext((>=)(x), buf, 1) + if !isnothing(i) + buf[i+1:end] = buf[i:end-1] + buf[i] = x + end + end + end + + for i = k+1:length(a) + fit_in(a[i]) + end + + buf[end] +end + +function quickselect_partition(a::AbstractVector, i, j) + x, y = i-1, j+1 + pivix = div(i+j, 2) + piv = a[pivix] + @show piv + while true + x += 1 + while a[x] < piv + x += 1 + end + y -= 1 + while a[y] > piv + y -= 1 + end + if x >= y + a[pivix], a[y] = a[y], a[pivix] + return y + end + if a[x] == piv + pivix = y + elseif a[y] == piv + pivix = x + end + a[x], a[y] = a[y], a[x] + end +end + +function quickselect(a::AbstractVector{<:Integer}, k=3, i=1, j=length(a)) + # Like quicksort, but keep sorting only half. + if i >= j + return i + end + pivix = quickselect_partition(a, i, j) + if pivix > k-1 + quickselect(a, k, i, pivix) + else + quickselect(a, k, pivix+1, j) + end +end