From bbf781f694f79bf4da7b8478bc2ee9b327e1adb5 Mon Sep 17 00:00:00 2001
From: Ray Miller <ray@1729.org.uk>
Date: Sat, 7 Dec 2024 13:14:58 +0000
Subject: [PATCH] Add priority-queue implementation

Heap-based implementaiton of a priority queue.
---
 algo/priority-queue.scm | 114 ++++++++++++++++++++++++++++++++++++++++
 1 file changed, 114 insertions(+)
 create mode 100644 algo/priority-queue.scm

diff --git a/algo/priority-queue.scm b/algo/priority-queue.scm
new file mode 100644
index 0000000..1f49aef
--- /dev/null
+++ b/algo/priority-queue.scm
@@ -0,0 +1,114 @@
+(define-module (algo priority-queue)
+  #:use-module (ice-9 match)
+  #:export (make-priority-queue
+            pq-length
+            pq-push!
+            pq-pop!
+            pq-peek
+            pq-empty?))
+
+;; The heap implementation is the based on Chapter 6: Heapsort of "Introduction
+;; to Algorithms" by Cormen, Leiserson, Rivest, and Stein.
+
+;; We store the heap in a vector with the first element storing the
+;; current heap size. This means elements of the heap are indexed starting
+;; from 1.
+
+(define (parent i)
+  (floor-quotient i 2))
+
+(define (left-child i)
+  (* 2 i))
+
+(define (right-child i)
+  (1+ (* 2 i)))
+
+(define (vector-swap! A i j)
+  (let ((tmp (vector-ref A i)))
+    (vector-set! A i (vector-ref A j))
+    (vector-set! A j tmp)))
+
+;; The sink function is called to maintain the heap property when the root node
+;; is removed and swapped with the last node in the heap. The new root node
+;; might need to "sink" down to maintain the heap ordering.
+;;
+;; priority is a function that returns the priority of an element.
+;; cmp will be > for a max heap or < for a min heap.
+(define (sink A i priority cmp)
+  (let ((heap-size (vector-ref A 0)))
+    (let loop ((i i))
+      (let ((l (left-child i))
+            (r (right-child i)))
+        (cond
+         ((and (<= r heap-size)
+               (cmp (priority (vector-ref A r))
+                    (priority (vector-ref A l))
+                    (priority (vector-ref A i))))
+          (vector-swap! A i r)
+          (loop r))
+         ((and (<= l heap-size)
+               (cmp (priority (vector-ref A l))
+                    (priority (vector-ref A i))))
+          (vector-swap! A i l)
+          (loop l)))))))
+
+;; The swim function is used to restore the heap property when a new node is
+;; added to the heap: it's initally added at the end of the heap with and
+;; may need to "swim" up the heap to maintain the order.
+(define (swim A i priority cmp)
+  (let loop ((i i))
+    (when (and (> (parent i) 0) (not (cmp (priority (vector-ref A (parent i)))
+                                          (priority (vector-ref A i)))))
+      (vector-swap! A i (parent i))
+      (loop (parent i)))))
+
+;; make-priority-queue constructs a new heap-based priority queue
+;; cmp should be > for a max-priority queue or < for a min-priority queue
+;; priority is a function that returns the priority of an element
+(define* (make-priority-queue #:key (cmp >) (priority identity) (initial-capacity 16))
+  (let ((A (make-vector initial-capacity)))
+    (vector-set! A 0 0)
+    (match-lambda*
+      (('dump)
+       A)
+      (('size)
+       (vector-ref A 0))
+      (('peek)
+       (let ((heap-size (vector-ref A 0)))
+         (if (< heap-size 1)
+             (error "heap underflow")
+             (vector-ref A 1))))
+      (('pop!)
+       (let ((heap-size (vector-ref A 0)))
+         (if (< heap-size 1)
+             (error "heap underflow")
+             (let ((result (vector-ref A 1)))
+               (vector-set! A 1 (vector-ref A heap-size))
+               (vector-set! A 0 (1- heap-size))
+               (sink A 1 priority cmp)
+               result))))
+      (('push! v)
+       (let ((heap-size (vector-ref A 0))
+             (capacity (vector-length A)))
+         (when (>= heap-size (1- capacity))
+           (let ((new-A (make-vector (* 2 capacity))))
+             (vector-copy! new-A 0 A)
+             (set! A new-A)))
+         (let ((heap-size (1+ heap-size)))
+           (vector-set! A 0 heap-size)
+           (vector-set! A heap-size v)
+           (swim A heap-size priority cmp)))))))
+
+(define (pq-length q)
+  (q 'size))
+
+(define (pq-push! q v)
+  (q 'push! v))
+
+(define (pq-pop! q)
+  (q 'pop!))
+
+(define (pq-peek q)
+  (q 'peek))
+
+(define pq-empty? (compose zero? pq-length))