Create a gist now

Instantly share code, notes, and snippets.

What would you like to do?
fbin.f : Fast equally-populated bins using the quickselect algorithm
module fbin
use iso_c_binding
implicit none
contains
subroutine eqpop_bin(X, Nb, Ntrl, Xout)
! ARG
integer, intent(in) :: Nb, Ntrl
real(c_double), intent(in) :: X(Ntrl)
integer(c_int8_t), intent(out) :: Xout(Ntrl)
! LOC
integer :: ks(Nb-1)
real(c_double) :: arr(Ntrl)
integer :: idx(Ntrl), i, j, left, right, startidx
! copy for in place manipulation
arr = X
! calculate kth elements required
do i=1,(Nb-1)
ks(i) = nint( (Ntrl*i) / real(Nb) )
end do
! fill index
do i=1,Ntrl
idx(i) = i
end do
left = 1
right = Ntrl
call selectkth(ks(1), arr, idx, left, right)
do i=2,(Nb-1)
left = ks(i-1) + 1;
call selectkth(ks(i), arr, idx, left, right)
end do
! now assign bin labels
startidx = 1
do i=1,(Nb-1)
do j=startidx,ks(i)
Xout(idx(j)) = i-1
end do
startidx = ks(i) + 1
end do
do i=startidx,Ntrl
Xout(idx(i)) = Nb - 1
end do
end subroutine eqpop_bin
subroutine selectkth(k, arr, idx, left_in, right_in)
! ARG
real(c_double), intent(inout) :: arr(:)
integer, intent(inout) :: idx(:)
integer, intent(in) :: k, left_in, right_in
! LOC
integer :: r, w, mididx, tmp_i, left, right
real(c_double) :: piv, tmp_r
left = left_in
right = right_in
! median of three approach to choose pivot
do while (left < right)
r = left
w = right
mididx = floor( (r+w) / 2.0 )
! sort with three conditional swaps
if( mididx > r ) then
if( arr(mididx) < arr(r) ) then
tmp_r = arr(mididx); arr(mididx) = arr(r); arr(r) = tmp_r
tmp_i = idx(mididx); idx(mididx) = idx(r); idx(r) = tmp_i
endif
if( arr(w) < arr(mididx) ) then
tmp_r = arr(mididx); arr(mididx) = arr(w); arr(w) = tmp_r
tmp_i = idx(mididx); idx(mididx) = idx(w); idx(w) = tmp_i
endif
if( arr(mididx) < arr(r) ) then
tmp_r = arr(mididx); arr(mididx) = arr(r); arr(r) = tmp_r
tmp_i = idx(mididx); idx(mididx) = idx(r); idx(r) = tmp_i
endif
endif
piv = arr(mididx)
do while (r < w)
if( arr(r) >= piv ) then
tmp_r = arr(w); arr(w) = arr(r); arr(r) = tmp_r
tmp_i = idx(w); idx(w) = idx(r); idx(r) = tmp_i
w = w - 1
else
r = r + 1
endif
end do
if( arr(r) > piv ) then
r = r - 1
endif
if( k <= r ) then
right = r;
else
left = r+1;
endif
end do
end subroutine selectkth
subroutine eqpop_bin_sorted(X, Nb, Ntrl, Xout)
! ARG
integer, intent(in) :: Nb, Ntrl
real(c_double), intent(in) :: X(Ntrl)
integer(c_int16_t), intent(out) :: Xout(Ntrl)
! LOC
integer :: ks(Nb-1)
integer :: i, j, startidx
! calculate kth elements required
do i=1,(Nb-1)
ks(i) = nint( (Ntrl*i) / real(Nb) )
end do
! assign bin labels
startidx = 1
do i=1,(Nb-1)
do j=startidx,ks(i)
Xout(j) = i-1
end do
startidx = ks(i) + 1
end do
do i=startidx,Ntrl
Xout(i) = Nb - 1
end do
end subroutine eqpop_bin_sorted
end module fbin
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment