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