堆排序以及堆结构的实现

roxillius

前一阵子在学python途中仔细读了一个heap库，觉得堆排序这个算法其实内容很丰富，于是自己就写了一个堆（二叉树）的模块。这个库实现了将一个数组整理为堆的功能(heapify)，另外排序功能是由其中n_largest这个函数实现，与完整的堆排序相比优点在于如果只需要寻找最大／小的n个元素的时候不需要将整个数组完全排序（这也是二叉树这个结构的优点）。目前只实现了降序排序，有兴趣的同学可以帮我添加升序排序。另外我写了两个heapify的函数，原因是为了比较两种不同整理思路的效率（两种不同的解释写在注释当中了，第二种思路是收到python heapq这个library的启发）。目前测来第一种heapify()比第二种heapify_alt()快不少，大致原因可能是我用节点数据结构相对简单的原因。
因为大家的水平有高有低，欢迎高手们给我拍砖，抓bug；也欢迎不懂的同学们提问题，我乐于解答／扫盲。

[Fortran] 纯文本查看 复制代码

module heap
  ! this module implements basic functions to realize a heap structure.
  implicit none
  save
  private 
  integer, parameter :: sp = selected_real_kind(6,37)
  integer, parameter :: i4b = selected_int_kind(9)
  public :: heapify
  public :: heapify_alt
!  public :: nsmallest
  public :: nlargest
  public :: sp
  public :: i4b
contains
  subroutine sift_down(pos,end_pos,arr)
    ! element array(pos) = temp is out of position, the function moves the temp element
    ! down "promotes" the descendent elements larger than array(pos), till array(pos)
    ! hits a leaf. The loop stops as well when array(pos) is larger than both daughters.
    ! this requires a supplementary comparison. 
    implicit none
    integer(i4b), intent(in) :: pos, end_pos
    real(sp), dimension(:), intent(inout) :: arr
    integer(i4b) :: j, jold
    real(sp) :: temp
    temp = arr(pos)
    jold = pos
    j = pos+pos  ! leftmost daughter of arr(pos) 
    do
       if (j>end_pos) exit  ! hitting the leaf
       if (j<end_pos) then  ! compare with the larger element between two daughters
          if (arr(j) < arr(j+1)) j = j + 1
       end if
       if (temp >= arr(j)) exit 
       arr(jold) = arr(j)
       jold = j
       j = j + j
    end do
    arr(jold) = temp
  end subroutine sift_down


  subroutine sift_up(pos,start_pos,arr)
    ! element arr(pos) is out of position, the heap is heap invariant at all
    ! positions <=start_pos. pos > startpos moves arr(pos) up and moves parent
    ! element down if smaller.
    implicit none
    real(sp), dimension(:), intent(inout) :: arr
    integer(i4b), intent(in) :: start_pos,pos
    integer(i4b) :: parent_pos, old_pos
    real(sp) :: temp
    old_pos = pos
    temp = arr(pos)
    do
       if (old_pos <= start_pos) exit
       parent_pos = old_pos/2
       if (arr(parent_pos) < temp) then
          arr(old_pos) = arr(parent_pos)
          old_pos = parent_pos
       else
          exit
       end if
    end do
    arr(old_pos) = temp
  end subroutine sift_up


  subroutine sift_down_alt(pos,end_pos,arr)
    ! alternative way of move the arr(pos) = temp down till to the leaf. the comparison
    ! between temp and the bigger daughter element is not made. As a result, the temp
    ! will reach the leaf no matter what. a sift down is performed afterwards to put
    ! the temp in proper position.
    implicit none
    integer(i4b), intent(in) :: pos, end_pos
    real(sp), dimension(:), intent(inout) :: arr
    real(sp) :: temp
    integer(i4b) :: daughter_pos, old_pos, start_pos
    temp = arr(pos)
    start_pos = pos
    old_pos = pos
    daughter_pos = pos+pos
    do
       if (daughter_pos > end_pos) exit  ! hitting the leaf
       if (daughter_pos < end_pos) then  ! if the element has got two daughters
          if (arr(daughter_pos) < arr(daughter_pos+1)) daughter_pos = daughter_pos + 1
       end if
       arr(old_pos) = arr(daughter_pos)
       old_pos = daughter_pos
       daughter_pos = daughter_pos + daughter_pos
    end do
    arr(old_pos) = temp
    call sift_up(old_pos,start_pos,arr)
  end subroutine sift_down_alt
    
  subroutine heapify(arr)
    implicit none
    real(sp), dimension(:), intent(inout) :: arr
    integer(i4b) :: i, n
    n = size(arr)
    do i = n/2, 1, -1
       call sift_down(i,n,arr)
    end do
  end subroutine heapify

  subroutine heapify_alt(arr)
    implicit none
    real(sp), dimension(:), intent(inout) :: arr
    integer(i4b) :: i, n
    n = size(arr)
    do i = n/2, 1, -1
       call sift_down_alt(i,n,arr)
    end do
  end subroutine heapify_alt

  function nlargest(arr, n) result(nlarge_arr)  ! return the first nth largest elements
    implicit none
    real(sp), dimension(:), intent(inout) :: arr
    integer(i4b), intent(in) :: n
    real(sp), dimension(n) :: nlarge_arr
    integer(i4b) :: i, j
    call heapify(arr)
    j = size(arr)
    do i = 1, n
       nlarge_arr(i) = arr(1)
       arr(1) = arr(j)
       arr(j) = nlarge_arr(i)
       j = j-1
       call sift_down(1,j,arr)
    end do
  end function nlargest
  
end module heap

测试文件已经添加于附件中。

pasuka

个人看法：
1、纯粹教学目的，Python、MATLAB这类脚本语言优于C、C++和Fortran这类编译型语言；
2、本站几年前有过相关讨论，lz不妨贴贴测试结果，Linux下面通过iso c binding调用qsort或许还略占优势
传送门
http://bbs.fcode.cn/thread-66-1-1.html