heap.F90

Go to the documentation of this file.

!! Copyright (C) 2021 S. Ohlmann
!!
!! This Source Code Form is subject to the terms of the Mozilla Public
!! License, v. 2.0. If a copy of the MPL was not distributed with this
!! file, You can obtain one at https://mozilla.org/MPL/2.0/.
!!
#include "global.h"
 
module heap_oct_m
  use debug_oct_m
  use global_oct_m
  use messages_oct_m
  use profiling_oct_m
  implicit none
  private
  public :: heap_t
 
  ! A min-heap type for implementing a k-way merge
  ! The min-heap is built on the first indices of the sorted arrays.
  ! It allows to always extract the minimum element of the sorted arrays
  ! and the heap structure is then maintained by swapping the indices and
  ! the corresponding sizes.
  ! To merge the sorted arrays, the minimum of the heap is extracted and
  ! replaced by the next element of that array; then the min-heap property
  ! is restored by down-sifting (minheapify).
  ! The time complexity of this merge is O(n log k) where n is the full size
  ! of the arrays and k is the number of arrays.
  type :: heap_t
    private
    integer(int64), pointer :: a(:) => null()  
    integer :: length
    integer, allocatable :: indices(:)
    integer, allocatable :: sizes(:)
  contains
    procedure :: init => heap_init
    procedure :: end => heap_end
    procedure :: merge => heap_merge
  end type heap_t
 
contains
  ! Initialize the heap
  subroutine heap_init(heap, list, sizes)
    class(heap_t),       intent(inout) :: heap
    integer(int64), target, intent(in)    :: list(:)
    integer,             intent(in)    :: sizes(:)
    integer :: i, j
 
    push_sub(heap_init)
    ! we keep a pointer to access the data; we do not need to modify it
    heap%a => list
 
    ! get the number of arrays to merge
    heap%length = ubound(sizes, dim=1)
    ! ignore empty parts
    do i = 1, ubound(sizes, dim=1)
      if (sizes(i) == 0) then
        heap%length = heap%length - 1
      end if
    end do
    safe_allocate(heap%sizes(heap%length))
    safe_allocate(heap%indices(heap%length))
    ! get the sizes of the non-empty arrays
    j = 1
    do i = 1, ubound(sizes, dim=1)
      if (sizes(i) == 0) cycle
      heap%sizes(j) = sizes(i)
      j = j + 1
    end do
    ! compute the starting indices
    heap%indices(1) = 1
    do i = 2, heap%length
      heap%indices(i) = heap%indices(i-1) + heap%sizes(i-1)
    end do
    ! sanity check
    if (sum(heap%sizes) /= ubound(list, dim=1)) then
      message(1) = "Error! Mismatch in size of array when initializing heap!"
      call messages_fatal(1)
    end if
 
    ! build the heap
    call build_minheap(heap)
    pop_sub(heap_init)
  end subroutine
 
  subroutine heap_end(heap)
    class(heap_t), intent(inout) :: heap
 
    push_sub(heap_end)
    nullify(heap%a)
    heap%length = 0
    safe_deallocate_a(heap%sizes)
    safe_deallocate_a(heap%indices)
    pop_sub(heap_end)
  end subroutine
 
  subroutine heap_merge(heap, merged, index_map)
    class(heap_t),     intent(inout) :: heap
    integer(int64),       intent(inout) :: merged(:)
    integer, optional, intent(inout) :: index_map(:)
 
    integer :: i
 
    push_sub(heap_merge)
 
    ! make sure the heap was initialized
    assert(associated(heap%a))
 
    if (sum(heap%sizes) /= ubound(merged, dim=1)) then
      message(1) = "Error! Mismatch in size of array when doing k-way merge!"
      call messages_fatal(1)
    end if
 
    do i = 1, sum(heap%sizes)
      ! extract the minimum element of the heap
      merged(i) = heap%a(heap%indices(1))
      if (present(index_map)) then
        index_map(i) = heap%indices(1)
      end if
      ! replace by the next element of that array and reduce its size
      heap%indices(1) = heap%indices(1) + 1
      heap%sizes(1) = heap%sizes(1) - 1
      if (heap%sizes(1) == 0) then
        ! if this array is already empty, swap it with the last one and
        ! reduce the size of the heap by one
        call swap(heap, 1_4, heap%length)
        heap%length = heap%length - 1
      end if
      ! down-sift (restore min-heap property)
      call minheapify(heap, 1_4)
    end do
    pop_sub(heap_merge)
  end subroutine
 
  ! build the minheap by downsifting all parent nodes
  subroutine build_minheap(heap)
    class(heap_t), intent(inout) :: heap
 
    integer :: i
 
    ! this goes bottom up for all parent nodes
    do i = heap%length/2, 1, -1
      call minheapify(heap, i)
    end do
  end subroutine
 
  ! downsifting from index i (restore minheap property)
  recursive subroutine minheapify(heap, i)
    class(heap_t),    intent(inout) :: heap
    integer,       intent(in)    :: i
 
    integer :: left, right, smallest
 
    left = 2*i
    right = 2*i + 1
    smallest = i
 
    ! get the smallest value of both children (if they are still in the heap)
    if (left <= heap%length) then
      if (is_smaller(heap, left, smallest)) then
        smallest = left
      end if
    end if
    if (right <= heap%length) then
      if (is_smaller(heap, right, smallest)) then
        smallest = right
      end if
    end if
    ! we only swap and downsift if one of the children is smaller
    if (smallest /= i) then
      call swap(heap, i, smallest)
      call minheapify(heap, smallest)
    end if
  end subroutine
 
  ! swap two of the sorted arrays in the heap
  subroutine swap(heap, i, j)
    class(heap_t),    intent(inout) :: heap
    integer,       intent(in)    :: i
    integer,       intent(in)    :: j
 
    ! swap index and associated size
    call swap_int(heap%indices, i, j)
    call swap_int(heap%sizes, i, j)
  end subroutine
 
  ! swap two integers
  subroutine swap_int(a, i, j)
    integer, intent(inout) :: a(:)
    integer, intent(in)    :: i
    integer, intent(in)    :: j
 
    integer :: tmp
 
    tmp = a(i)
    a(i) = a(j)
    a(j) = tmp
  end subroutine
 
  ! compare the current value of two arrays in the heap
  logical function is_smaller(heap, i, j)
    class(heap_t),    intent(inout) :: heap
    integer,       intent(in)    :: i
    integer,       intent(in)    :: j
 
    is_smaller = heap%a(heap%indices(i)) < heap%a(heap%indices(j))
  end function
end module heap_oct_m