min_heap.h 3.3 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134
  1. /* SPDX-License-Identifier: GPL-2.0 */
  2. #ifndef _LINUX_MIN_HEAP_H
  3. #define _LINUX_MIN_HEAP_H
  4. #include <linux/bug.h>
  5. #include <linux/string.h>
  6. #include <linux/types.h>
  7. /**
  8. * struct min_heap - Data structure to hold a min-heap.
  9. * @data: Start of array holding the heap elements.
  10. * @nr: Number of elements currently in the heap.
  11. * @size: Maximum number of elements that can be held in current storage.
  12. */
  13. struct min_heap {
  14. void *data;
  15. int nr;
  16. int size;
  17. };
  18. /**
  19. * struct min_heap_callbacks - Data/functions to customise the min_heap.
  20. * @elem_size: The nr of each element in bytes.
  21. * @less: Partial order function for this heap.
  22. * @swp: Swap elements function.
  23. */
  24. struct min_heap_callbacks {
  25. int elem_size;
  26. bool (*less)(const void *lhs, const void *rhs);
  27. void (*swp)(void *lhs, void *rhs);
  28. };
  29. /* Sift the element at pos down the heap. */
  30. static __always_inline
  31. void min_heapify(struct min_heap *heap, int pos,
  32. const struct min_heap_callbacks *func)
  33. {
  34. void *left, *right, *parent, *smallest;
  35. void *data = heap->data;
  36. for (;;) {
  37. if (pos * 2 + 1 >= heap->nr)
  38. break;
  39. left = data + ((pos * 2 + 1) * func->elem_size);
  40. parent = data + (pos * func->elem_size);
  41. smallest = parent;
  42. if (func->less(left, smallest))
  43. smallest = left;
  44. if (pos * 2 + 2 < heap->nr) {
  45. right = data + ((pos * 2 + 2) * func->elem_size);
  46. if (func->less(right, smallest))
  47. smallest = right;
  48. }
  49. if (smallest == parent)
  50. break;
  51. func->swp(smallest, parent);
  52. if (smallest == left)
  53. pos = (pos * 2) + 1;
  54. else
  55. pos = (pos * 2) + 2;
  56. }
  57. }
  58. /* Floyd's approach to heapification that is O(nr). */
  59. static __always_inline
  60. void min_heapify_all(struct min_heap *heap,
  61. const struct min_heap_callbacks *func)
  62. {
  63. int i;
  64. for (i = heap->nr / 2; i >= 0; i--)
  65. min_heapify(heap, i, func);
  66. }
  67. /* Remove minimum element from the heap, O(log2(nr)). */
  68. static __always_inline
  69. void min_heap_pop(struct min_heap *heap,
  70. const struct min_heap_callbacks *func)
  71. {
  72. void *data = heap->data;
  73. if (WARN_ONCE(heap->nr <= 0, "Popping an empty heap"))
  74. return;
  75. /* Place last element at the root (position 0) and then sift down. */
  76. heap->nr--;
  77. memcpy(data, data + (heap->nr * func->elem_size), func->elem_size);
  78. min_heapify(heap, 0, func);
  79. }
  80. /*
  81. * Remove the minimum element and then push the given element. The
  82. * implementation performs 1 sift (O(log2(nr))) and is therefore more
  83. * efficient than a pop followed by a push that does 2.
  84. */
  85. static __always_inline
  86. void min_heap_pop_push(struct min_heap *heap,
  87. const void *element,
  88. const struct min_heap_callbacks *func)
  89. {
  90. memcpy(heap->data, element, func->elem_size);
  91. min_heapify(heap, 0, func);
  92. }
  93. /* Push an element on to the heap, O(log2(nr)). */
  94. static __always_inline
  95. void min_heap_push(struct min_heap *heap, const void *element,
  96. const struct min_heap_callbacks *func)
  97. {
  98. void *data = heap->data;
  99. void *child, *parent;
  100. int pos;
  101. if (WARN_ONCE(heap->nr >= heap->size, "Pushing on a full heap"))
  102. return;
  103. /* Place at the end of data. */
  104. pos = heap->nr;
  105. memcpy(data + (pos * func->elem_size), element, func->elem_size);
  106. heap->nr++;
  107. /* Sift child at pos up. */
  108. for (; pos > 0; pos = (pos - 1) / 2) {
  109. child = data + (pos * func->elem_size);
  110. parent = data + ((pos - 1) / 2) * func->elem_size;
  111. if (func->less(parent, child))
  112. break;
  113. func->swp(parent, child);
  114. }
  115. }
  116. #endif /* _LINUX_MIN_HEAP_H */