Add priority-queue implementation

Heap-based implementaiton of a priority queue.
This commit is contained in:
Ray Miller 2024-12-07 13:14:58 +00:00
parent c6e8b71198
commit bbf781f694
Signed by: ray
GPG key ID: 043F786C4CD681B8

114
algo/priority-queue.scm Normal file
View file

@ -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))