cp_min_heap.F

Go to the documentation of this file.

!--------------------------------------------------------------------------------------------------!
!   CP2K: A general program to perform molecular dynamics simulations                              !
!   Copyright 2000-2024 CP2K developers group <https://cp2k.org>                                   !
!                                                                                                  !
!   SPDX-License-Identifier: GPL-2.0-or-later                                                      !
!--------------------------------------------------------------------------------------------------!
 
MODULE cp_min_heap
   USE kinds,                           ONLY: int_4,&
                                              int_8
#include "../base/base_uses.f90"
 
   IMPLICIT NONE
   PRIVATE
 
   PUBLIC :: cp_heap_type, keyt, valt
   PUBLIC :: cp_heap_pop, cp_heap_reset_node, cp_heap_fill
   PUBLIC :: cp_heap_new, cp_heap_release
   PUBLIC :: cp_heap_get_first, cp_heap_reset_first
 
   ! Sets the types
   INTEGER, PARAMETER         :: keyt = int_4
   INTEGER, PARAMETER         :: valt = int_8
 
   TYPE cp_heap_node
      INTEGER(KIND=keyt) :: key = -1_keyt
      INTEGER(KIND=valt) :: value = -1_valt
   END TYPE cp_heap_node
 
   TYPE cp_heap_node_e
      TYPE(cp_heap_node) :: node = cp_heap_node()
   END TYPE cp_heap_node_e
 
   TYPE cp_heap_type
      INTEGER :: n = -1
      INTEGER, DIMENSION(:), POINTER           :: index => null()
      TYPE(cp_heap_node_e), DIMENSION(:), POINTER :: nodes => null()
   END TYPE cp_heap_type
 
CONTAINS
 
   ! Lookup functions
 
! **************************************************************************************************
!> \brief ...
!> \param n ...
!> \return ...
! **************************************************************************************************
   ELEMENTAL FUNCTION get_parent(n) RESULT(parent)
      INTEGER, INTENT(IN)                                :: n
      INTEGER                                            :: parent
 
      parent = int(n/2)
   END FUNCTION get_parent
 
! **************************************************************************************************
!> \brief ...
!> \param n ...
!> \return ...
! **************************************************************************************************
   ELEMENTAL FUNCTION get_left_child(n) RESULT(child)
      INTEGER, INTENT(IN)                                :: n
      INTEGER                                            :: child
 
      child = 2*n
   END FUNCTION get_left_child
 
! **************************************************************************************************
!> \brief ...
!> \param n ...
!> \return ...
! **************************************************************************************************
   ELEMENTAL FUNCTION get_right_child(n) RESULT(child)
      INTEGER, INTENT(IN)                                :: n
      INTEGER                                            :: child
 
      child = 2*n + 1
   END FUNCTION get_right_child
 
! **************************************************************************************************
!> \brief ...
!> \param heap ...
!> \param n ...
!> \return ...
! **************************************************************************************************
   ELEMENTAL FUNCTION get_value(heap, n) RESULT(value)
      TYPE(cp_heap_type), INTENT(IN)                     :: heap
      INTEGER, INTENT(IN)                                :: n
      INTEGER(KIND=valt)                                 :: value
 
      value = heap%nodes(n)%node%value
   END FUNCTION get_value
 
! **************************************************************************************************
!> \brief ...
!> \param heap ...
!> \param key ...
!> \return ...
! **************************************************************************************************
   ELEMENTAL FUNCTION get_value_by_key(heap, key) RESULT(value)
      TYPE(cp_heap_type), INTENT(IN)                     :: heap
      INTEGER(KIND=keyt), INTENT(IN)                     :: key
      INTEGER(KIND=valt)                                 :: value
 
      INTEGER                                            :: n
 
      n = heap%index(key)
      value = get_value(heap, n)
   END FUNCTION get_value_by_key
 
   ! Initialization functions
 
! **************************************************************************************************
!> \brief ...
!> \param heap ...
!> \param n ...
! **************************************************************************************************
   SUBROUTINE cp_heap_new(heap, n)
      TYPE(cp_heap_type), INTENT(OUT)                    :: heap
      INTEGER, INTENT(IN)                                :: n
 
      heap%n = n
      ALLOCATE (heap%index(n))
      ALLOCATE (heap%nodes(n))
   END SUBROUTINE cp_heap_new
 
! **************************************************************************************************
!> \brief ...
!> \param heap ...
! **************************************************************************************************
   SUBROUTINE cp_heap_release(heap)
      TYPE(cp_heap_type), INTENT(INOUT)                  :: heap
 
      DEALLOCATE (heap%index)
      DEALLOCATE (heap%nodes)
      heap%n = 0
   END SUBROUTINE cp_heap_release
 
! **************************************************************************************************
!> \brief Fill heap with given values
!> \param heap ...
!> \param values ...
! **************************************************************************************************
   SUBROUTINE cp_heap_fill(heap, values)
      TYPE(cp_heap_type), INTENT(INOUT)                  :: heap
      INTEGER(KIND=valt), DIMENSION(:), INTENT(IN)       :: values
 
      INTEGER                                            :: first, i, n
 
      n = SIZE(values)
      cpassert(heap%n >= n)
 
      DO i = 1, n
         heap%index(i) = i
         heap%nodes(i)%node%key = i
         heap%nodes(i)%node%value = values(i)
      END DO
      ! Sort from the last full subtree
      first = get_parent(n)
      DO i = first, 1, -1
         CALL bubble_down(heap, i)
      END DO
 
   END SUBROUTINE cp_heap_fill
 
! **************************************************************************************************
!> \brief Returns the first heap element without removing it.
!> \param heap ...
!> \param key ...
!> \param value ...
!> \param found ...
! **************************************************************************************************
   SUBROUTINE cp_heap_get_first(heap, key, value, found)
      TYPE(cp_heap_type), INTENT(INOUT)                  :: heap
      INTEGER(KIND=keyt), INTENT(OUT)                    :: key
      INTEGER(KIND=valt), INTENT(OUT)                    :: value
      LOGICAL, INTENT(OUT)                               :: found
 
      IF (heap%n .LT. 1) THEN
         found = .false.
      ELSE
         found = .true.
         key = heap%nodes(1)%node%key
         value = heap%nodes(1)%node%value
      END IF
   END SUBROUTINE cp_heap_get_first
 
! **************************************************************************************************
!> \brief Returns and removes the first heap element and rebalances
!>        the heap.
!> \param heap ...
!> \param key ...
!> \param value ...
!> \param found ...
! **************************************************************************************************
   SUBROUTINE cp_heap_pop(heap, key, value, found)
      TYPE(cp_heap_type), INTENT(INOUT)                  :: heap
      INTEGER(KIND=keyt), INTENT(OUT)                    :: key
      INTEGER(KIND=valt), INTENT(OUT)                    :: value
      LOGICAL, INTENT(OUT)                               :: found
 
!
 
      CALL cp_heap_get_first(heap, key, value, found)
      IF (found) THEN
         IF (heap%n .GT. 1) THEN
            CALL cp_heap_copy_node(heap, 1, heap%n)
            heap%n = heap%n - 1
            CALL bubble_down(heap, 1)
         ELSE
            heap%n = heap%n - 1
         END IF
      END IF
   END SUBROUTINE cp_heap_pop
 
! **************************************************************************************************
!> \brief Changes the value of the heap element with given key and
!>        rebalances the heap.
!> \param heap ...
!> \param key ...
!> \param value ...
! **************************************************************************************************
   SUBROUTINE cp_heap_reset_node(heap, key, value)
      TYPE(cp_heap_type), INTENT(INOUT)                  :: heap
      INTEGER(KIND=keyt), INTENT(IN)                     :: key
      INTEGER(KIND=valt), INTENT(IN)                     :: value
 
      INTEGER                                            :: n, new_pos
 
      cpassert(heap%n > 0)
 
      n = heap%index(key)
      cpassert(heap%nodes(n)%node%key == key)
      heap%nodes(n)%node%value = value
      CALL bubble_up(heap, n, new_pos)
      CALL bubble_down(heap, new_pos)
   END SUBROUTINE cp_heap_reset_node
 
! **************************************************************************************************
!> \brief Changes the value of the minimum heap element and rebalances the heap.
!> \param heap ...
!> \param value ...
! **************************************************************************************************
   SUBROUTINE cp_heap_reset_first(heap, value)
      TYPE(cp_heap_type), INTENT(INOUT)                  :: heap
      INTEGER(KIND=valt), INTENT(IN)                     :: value
 
      cpassert(heap%n > 0)
      heap%nodes(1)%node%value = value
      CALL bubble_down(heap, 1)
   END SUBROUTINE cp_heap_reset_first
 
! **************************************************************************************************
!> \brief ...
!> \param heap ...
!> \param e1 ...
!> \param e2 ...
! **************************************************************************************************
   PURE SUBROUTINE cp_heap_swap(heap, e1, e2)
      TYPE(cp_heap_type), INTENT(INOUT)                  :: heap
      INTEGER, INTENT(IN)                                :: e1, e2
 
      INTEGER(KIND=keyt)                                 :: key1, key2
      TYPE(cp_heap_node)                                 :: tmp_node
 
      key1 = heap%nodes(e1)%node%key
      key2 = heap%nodes(e2)%node%key
 
      tmp_node = heap%nodes(e1)%node
      heap%nodes(e1)%node = heap%nodes(e2)%node
      heap%nodes(e2)%node = tmp_node
 
      heap%index(key1) = e2
      heap%index(key2) = e1
   END SUBROUTINE cp_heap_swap
 
! **************************************************************************************************
!> \brief Sets node e1 to e2
!> \param heap ...
!> \param e1 ...
!> \param e2 ...
! **************************************************************************************************
   PURE SUBROUTINE cp_heap_copy_node(heap, e1, e2)
      TYPE(cp_heap_type), INTENT(INOUT)                  :: heap
      INTEGER, INTENT(IN)                                :: e1, e2
 
      INTEGER(KIND=keyt)                                 :: key1, key2
 
      key1 = heap%nodes(e1)%node%key
      key2 = heap%nodes(e2)%node%key
 
      heap%nodes(e1)%node = heap%nodes(e2)%node
 
      heap%index(key1) = 0
      heap%index(key2) = e1
   END SUBROUTINE cp_heap_copy_node
 
! **************************************************************************************************
!> \brief Balances a heap by bubbling down from the given element.
!> \param heap ...
!> \param first ...
! **************************************************************************************************
   SUBROUTINE bubble_down(heap, first)
      TYPE(cp_heap_type), INTENT(INOUT)                  :: heap
      INTEGER, INTENT(IN)                                :: first
 
      INTEGER                                            :: e, left_child, right_child, smallest
      INTEGER(kind=valt)                                 :: left_child_value, min_value, &
                                                            right_child_value
      LOGICAL                                            :: all_done
 
!
      cpassert(0 < first .AND. first <= heap%n)
 
      e = first
      all_done = .false.
      ! Check whether we are finished, i.e,. whether the element to
      ! bubble down is childless.
      DO WHILE (e .LE. get_parent(heap%n) .AND. .NOT. all_done)
         ! Determines which node (current, left, or right child) has the
         ! smallest value.
         smallest = e
         min_value = get_value(heap, e)
         left_child = get_left_child(e)
         IF (left_child .LE. heap%n) THEN
            left_child_value = get_value(heap, left_child)
            IF (left_child_value .LT. min_value) THEN
               min_value = left_child_value
               smallest = left_child
            END IF
         END IF
         right_child = left_child + 1
         IF (right_child .LE. heap%n) THEN
            right_child_value = get_value(heap, right_child)
            IF (right_child_value .LT. min_value) THEN
               min_value = right_child_value
               smallest = right_child
            END IF
         END IF
         !
         CALL cp_heap_swap(heap, e, smallest)
         IF (smallest .EQ. e) THEN
            all_done = .true.
         ELSE
            e = smallest
         END IF
      END DO
   END SUBROUTINE bubble_down
 
! **************************************************************************************************
!> \brief Balances a heap by bubbling up from the given element.
!> \param heap ...
!> \param first ...
!> \param new_pos ...
! **************************************************************************************************
   SUBROUTINE bubble_up(heap, first, new_pos)
      TYPE(cp_heap_type), INTENT(INOUT)                  :: heap
      INTEGER, INTENT(IN)                                :: first
      INTEGER, INTENT(OUT)                               :: new_pos
 
      INTEGER                                            :: e, parent
      INTEGER(kind=valt)                                 :: my_value, parent_value
      LOGICAL                                            :: all_done
 
      cpassert(0 < first .AND. first <= heap%n)
 
      e = first
      all_done = .false.
      IF (e .GT. 1) THEN
         my_value = get_value(heap, e)
      END IF
      ! Check whether we are finished, i.e,. whether the element to
      ! bubble up is an orphan.
      new_pos = e
      DO WHILE (e .GT. 1 .AND. .NOT. all_done)
         ! Switches the parent and the current element if the current
         ! element's value is greater than the parent's value.
         parent = get_parent(e)
         parent_value = get_value(heap, parent)
         IF (my_value .LT. parent_value) THEN
            CALL cp_heap_swap(heap, e, parent)
            e = parent
         ELSE
            all_done = .true.
         END IF
      END DO
      new_pos = e
   END SUBROUTINE bubble_up
 
END MODULE cp_min_heap