priority_queue.cpp 2.4 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596
  1. template <typename T>
  2. struct PriorityQueue {
  3. Array<T> queue;
  4. int (* cmp) (T *q, isize i, isize j);
  5. void (* swap)(T *q, isize i, isize j);
  6. };
  7. template <typename T>
  8. gb_internal bool priority_queue_shift_down(PriorityQueue<T> *pq, isize i0, isize n) {
  9. // O(n log n)
  10. isize i = i0;
  11. isize j, j1, j2;
  12. if (0 > i || i > n) return false;
  13. for (;;) {
  14. j1 = 2*i + 1;
  15. if (0 > j1 || j1 >= n) break;
  16. j = j1;
  17. j2 = j1 + 1;
  18. if (j2 < n && pq->cmp(&pq->queue[0], j2, j1) < 0) {
  19. j = j2;
  20. }
  21. if (pq->cmp(&pq->queue[0], j, i) >= 0) break;
  22. pq->swap(&pq->queue[0], i, j);
  23. i = j;
  24. }
  25. return i > i0;
  26. }
  27. template <typename T>
  28. gb_internal void priority_queue_shift_up(PriorityQueue<T> *pq, isize j) {
  29. while (0 <= j && j < pq->queue.count) {
  30. isize i = (j-1)/2;
  31. if (i == j || pq->cmp(&pq->queue[0], j, i) >= 0) {
  32. break;
  33. }
  34. pq->swap(&pq->queue[0], i, j);
  35. j = i;
  36. }
  37. }
  38. // NOTE(bill): When an element at index `i0` has changed its value, this will fix the
  39. // the heap ordering. This using a basic "heapsort" with shift up and a shift down parts.
  40. template <typename T>
  41. gb_internal void priority_queue_fix(PriorityQueue<T> *pq, isize i) {
  42. if (!priority_queue_shift_down(pq, i, pq->queue.count)) {
  43. priority_queue_shift_up(pq, i);
  44. }
  45. }
  46. template <typename T>
  47. gb_internal void priority_queue_push(PriorityQueue<T> *pq, T const &value) {
  48. array_add(&pq->queue, value);
  49. priority_queue_shift_up(pq, pq->queue.count-1);
  50. }
  51. template <typename T>
  52. gb_internal T priority_queue_pop(PriorityQueue<T> *pq) {
  53. GB_ASSERT(pq->queue.count > 0);
  54. isize n = pq->queue.count - 1;
  55. pq->swap(&pq->queue[0], 0, n);
  56. priority_queue_shift_down(pq, 0, n);
  57. return array_pop(&pq->queue);
  58. }
  59. template <typename T>
  60. gb_internal T priority_queue_remove(PriorityQueue<T> *pq, isize i) {
  61. GB_ASSERT(0 <= i && i < pq->queue.count);
  62. isize n = pq->queue.count - 1;
  63. if (n != i) {
  64. pq->swap(&pq->queue[0], i, n);
  65. priority_queue_shift_down(pq, i, n);
  66. priority_queue_shift_up(pq, i);
  67. }
  68. return array_pop(&pq->queue);
  69. }
  70. template <typename T>
  71. gb_internal PriorityQueue<T> priority_queue_create(Array<T> queue,
  72. int (* cmp) (T *q, isize i, isize j),
  73. void (* swap)(T *q, isize i, isize j)) {
  74. PriorityQueue<T> pq = {};
  75. pq.queue = queue;
  76. pq.cmp = cmp;
  77. pq.swap = swap;
  78. isize n = pq.queue.count;
  79. for (isize i = n/2 - 1; i >= 0; i--) {
  80. priority_queue_shift_down(&pq, i, n);
  81. }
  82. return pq;
  83. }