fbin.f : Fast equally-populated bins using the quickselect algorithm

fbin.f

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