haskeem-0.7.5: heap.scm

; Copyright 2009 Uwe Hollerbach <uh@alumni.caltech.edu>
; $Id: heap.scm,v 1.2 2009-06-29 00:24:24 uwe Exp $
; BSD3

; Priority queues implemented as heaps

; What's a heap? It's a vector containing another vector:
;
;	#(pred? cur-size #(array))
;
; where pred? is a predicate used to maintain the heap structure, cur-size
; is the number of items currently in the heap, and array is the actual
; heap.

; Create a new heap which will use the given predicate for comparisons, and
; set it to initially have the given size. Vectors can be resized, so it's
; not a big deal to get the size wrong.

(define (heap-new pred? size)
  (let ((h (make-vector 3))
	(s (max size 1)))
    (vector-set! h 0 pred?)
    (vector-set! h 1 0)
    (vector-set! h 2 (make-vector s))))

; Insert an item into a heap, maintaining the heap condition;
; Return the modified heap

(define (heap-insert heap item)
  (let* ((pred? (vector-ref heap 0))
	 (k (vector-ref heap 1))
	 (a (vector-ref heap 2))
	 (n (vector-length a))
	 (kp #f)
	 (parent (lambda (j) (quotient (- j 1) 2))))
; make sure there's space for insertion
    (when (= k n)
	  (set! n (* n 2))
	  (vector-resize! a n))
; insert item at end of heap
    (vector-set! a k item)
    (vector-set! heap 1 (+ k 1))
; restore heap condition
    (do ((kp (parent k) (parent k)))
	((or (zero? k) (pred? (vector-ref a kp) item)) #t)
      (vector-set! a k (vector-ref a kp))
      (set! k kp))
    (vector-set! a k item)
; save modified array, which also returns heap to caller
    (vector-set! heap 2 a)))

; Remove the highest-priority item from the heap, maintaining the heap
; condition; return a list of the highest-priority item and the new heap

(define (heap-remove heap)
  (let* ((pred? (vector-ref heap 0))
	 (k (vector-ref heap 1))
	 (a (vector-ref heap 2))
	 (kp 0)
	 (ret (if (zero? k) (raise "heap exhausted!") (vector-ref a kp)))
	 (kc #f)
	 (v #f)
	 (child (lambda (j) (+ 1 (* 2 j)))))
    (set! k (- k 1))
    (vector-set! heap 1 k)
    (set! v (vector-ref a k))
    (vector-set! a kp v)
    (set! k (- k 1))
    (set! kc (child kp))
    (when (and (< kc k) (pred? (vector-ref a (+ kc 1))
			       (vector-ref a kc)))
	  (set! kc (+ kc 1)))
    (while (and (<= (child kp) k) (pred? (vector-ref a kc) v))
	   (vector-set! a kp (vector-ref a kc))
	   (set! kp kc)
	   (set! kc (child kp))
	   (when (and (< kc k) (pred? (vector-ref a (+ kc 1))
				      (vector-ref a kc)))
		 (set! kc (+ kc 1))))
    (vector-set! a kp v)
    (vector-set! heap 2 a)
    (list ret heap)))

; Insert an item into a heap, updating the heap in place

(defmacro (heap-insert! heap item)
  `(set! ,heap (heap-insert ,heap ,item)))

; Remove the highest-priority item from the heap, updating the heap
; in place, and return the highest-priority item

(defmacro (heap-remove! heap)
  (let ((l (new-symbol)))
    `(let ((,l (heap-remove ,heap)))
       (set! ,heap (cadr ,l))
       (car ,l))))

;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;
; debugging stuff

(define (display-heap heap label)
  (letrec* ((k (vector-ref heap 1))
	    (a (vector-ref heap 2))
	    (i 0)
	    (target 1)
	    (nrows (ilog 2 k))
	    (sp-count (lambda (x)
			(if (zero? x) 1 (+ 2 (* 2 (sp-count (- x 1)))))))
	    (nsp (sp-count nrows)))
    (write-string label #\newline)
    (do ((i 0 (+ i 1)))
	((>= i k) #t)
      (display (vector-ref a i))
      (if (= (+ i 1) target)
	  (begin (newline)
		 (set! target (+ 1 (* 2 target)))
		 (set! nsp (/ (- nsp 2) 2)))
	  (apply write-string (replicate #\space nsp))))
    (newline)))

; Check that the specified heap condition is met at the given index:
; compare value against its parent, which must satisfy (pred? parent child).
; Due to possibility of equality, check for existence of failure condition
; (pred? child parent) instead of success condition (pred? parent child).

(define (heap-check-one heap cur)
  (let ((pred? (vector-ref heap 0))
	(k (vector-ref heap 1))
	(a (vector-ref heap 2))
	(parent (lambda (j) (quotient (- j 1) 2))))
    (cond ((or (negative? cur) (>= cur k))
	   (raise "requested index is out of bounds!"))
	  ((zero? cur) (write-string "heap root\n"))
	  (else (let* ((cp (parent cur))
		       (vc (vector-ref a cur))
		       (vp (vector-ref a cp)))
		  (when (pred? vc vp)
			(write-string "heap condition not satisfied between ")
			(display vp)
			(write-string " at " (number->string cp) " and ")
			(display vc)
			(write-string " at " (number->string cur)
				      #\newline)))))))

; Check that the heap condition is met at all indices

(define (heap-check-all heap)
  (let ((k (vector-ref heap 1))
	(i 1))
    (do ((i 1 (+ i 1)))
	((>= i k) #t)
      (heap-check-one heap i))))

(define (test pred . vals)
  (let ((hp (heap-new pred 10))
	(i #f))
    (until (null? vals)
	   (heap-insert! hp (car vals))
	   (set! vals (cdr vals)))
    (heap-check-all hp)
    (while (positive? (vector-ref hp 1))
	   (set! i (heap-remove! hp))
	   (heap-check-all hp)
	   (write-string "removed ")
	   (display i)
	   (newline))))