+;;; -*-lisp-*-
+;;;
+;;; $Id$
+;;;
+;;; Heap data structure; useful for priority queues and suchlike
+;;;
+;;; (c) 2006 Straylight/Edgeware
+;;;
+
+;;;----- Licensing notice ---------------------------------------------------
+;;;
+;;; This program is free software; you can redistribute it and/or modify
+;;; it under the terms of the GNU General Public License as published by
+;;; the Free Software Foundation; either version 2 of the License, or
+;;; (at your option) any later version.
+;;;
+;;; This program is distributed in the hope that it will be useful,
+;;; but WITHOUT ANY WARRANTY; without even the implied warranty of
+;;; MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
+;;; GNU General Public License for more details.
+;;;
+;;; You should have received a copy of the GNU General Public License
+;;; along with this program; if not, write to the Free Software Foundation,
+;;; Inc., 59 Temple Place - Suite 330, Boston, MA 02111-1307, USA.
+
+(defpackage #:heap
+ (:use #:common-lisp)
+ (:export #:make-heap #:heap-count #:heap-empty-p
+ #:heap-insert #:heap-head #:heap-remove
+ #:heap-sort))
+(in-package #:heap)
+
+;;;--------------------------------------------------------------------------
+;;; Useful indexing functions.
+
+(declaim (inline parent left-child right-child))
+(deftype index () '(and unsigned-byte fixnum))
+(defun parent (i)
+ (declare (type index i))
+ (the index (floor (- i 1) 2)))
+(defun left-child (i)
+ (declare (type index i))
+ (the index (+ (* 2 i) 1)))
+(defun right-child (i)
+ (declare (type index i))
+ (the index (+ (* 2 i) 2)))
+
+;;;--------------------------------------------------------------------------
+;;; Low-level heap operations.
+
+(defun upheap (v key cmp n x)
+ "Insert the element X in the highest place possible in the heap."
+ (declare (type vector v)
+ (type function key cmp)
+ (type index n))
+ (let ((i n) (xk (funcall key x)))
+ (loop (when (zerop i) (return))
+ (let* ((j (parent i))
+ (y (aref v j)))
+ (when (funcall cmp (funcall key y) xk) (return))
+ (setf (aref v i) y
+ i j)))
+ (setf (aref v i) x)))
+
+(defun downheap (v key cmp n x)
+ "Insert the element X in the lowest place possible in the heap."
+ (declare (type vector v)
+ (type function key cmp)
+ (type index n))
+ (let ((i 0) (xk (funcall key x)))
+ (loop (let ((l (left-child i))
+ (r (right-child i)))
+ (when (>= l n) (return))
+ (multiple-value-bind
+ (j y yk)
+ (let* ((y (aref v l))
+ (yk (funcall key y)))
+ (if (= r n)
+ (values l y yk)
+ (let* ((z (aref v r))
+ (zk (funcall key z)))
+ (if (funcall cmp yk zk)
+ (values l y yk)
+ (values r z zk)))))
+ (when (funcall cmp xk yk)
+ (return))
+ (setf (aref v i) y
+ i j))))
+ (setf (aref v i) x)))
+
+(defun check (v key cmp n)
+ "Verify the heap invariant on the heap."
+ (declare (type vector v)
+ (type function key cmp)
+ (type index n))
+ (dotimes (i n)
+ (let* ((item (aref v i))
+ (item-key (funcall key item))
+ (l (left-child i))
+ (r (right-child i)))
+ (when (< l n)
+ (let ((left-item (aref v l)))
+ (assert (funcall cmp item-key (funcall key left-item))))
+ (when (< r n)
+ (let ((right-item (aref v r)))
+ (assert (funcall cmp item-key (funcall key right-item)))))))))
+
+;;;--------------------------------------------------------------------------
+;;; High-level heap things
+
+(defstruct (heap (:predicate heapp) (:constructor %make-heap))
+ "Data structure for a heap."
+ (v (make-array 16) :type vector)
+ (n 0 :type index)
+ (key #'identity :type function)
+ (compare #'<= :type function))
+
+(defun make-heap
+ (&key (compare #'<=) (key #'identity)
+ (type 't) (init-size 16) (contents nil contentsp))
+ "Return a new heap.
+
+ COMPARE is a partial-order predicate: (COMPARE X Y) should return true if
+ X <= Y in some order.
+
+ The TYPE is the element type of the heap.
+
+ INIT-SIZE is the initial allocation for the heap; the heap will grow
+ automatically if necessary, so this isn't a big deal. This is only a
+ hint; make-heap may ignore it completely.
+
+ KEY is a function to extract the key from an element. The default is to
+ use the item unmolested.
+
+ CONTENTS is the initial contents of the heap. If omitted, the heap is
+ initially empty."
+ (let ((n (if contentsp (length contents) 0)))
+ (loop while (< init-size n)
+ do (setf init-size (ash init-size 1)))
+ (let ((v (make-array init-size :element-type type)))
+ (when contentsp
+ (reduce (lambda (i item)
+ (upheap v key compare i item)
+ (1+ i))
+ contents
+ :initial-value 0))
+ (%make-heap :compare compare :key key :n n :v v))))
+
+(defun heap-count (heap)
+ "Return the number of elements in HEAP."
+ (declare (type heap heap))
+ (heap-n heap))
+
+(defun heap-empty-p (heap)
+ "True if HEAP is empty."
+ (declare (type heap heap))
+ (zerop (heap-count heap)))
+
+(defun heap-insert (heap item)
+ "Insert ITEM into the HEAP."
+ (declare (type heap heap))
+ (let* ((v (heap-v heap))
+ (n (heap-n heap))
+ (sz (array-dimension v 0)))
+ (when (= n sz)
+ (setf v (adjust-array v (* 2 n))
+ (heap-v heap) v))
+ (upheap v (heap-key heap) (heap-compare heap) n item)
+ (setf (heap-n heap) (1+ n))))
+
+(defun heap-head (heap)
+ "Peep at the head item on HEAP."
+ (declare (type heap heap))
+ (assert (not (heap-empty-p heap)))
+ (aref (heap-v heap) 0))
+
+(defun heap-remove (heap)
+ "Remove the head item from HEAP and return it."
+ (declare (type heap heap))
+ (assert (not (heap-empty-p heap)))
+ (let ((v (heap-v heap))
+ (n (1- (heap-n heap))))
+ (prog1 (aref v 0)
+ (setf (heap-n heap) n)
+ (downheap v (heap-key heap) (heap-compare heap) n (aref v n)))))
+
+(defun heap-sort (items compare &key (key #'identity))
+ "Return the ITEMS, least-first, as sorted by the ordering COMPARE."
+ (let ((heap (make-heap :compare compare :contents items :key key)))
+ (loop repeat (heap-n heap)
+ collect (heap-remove heap))))
+
+;;;----- That's all, folks --------------------------------------------------
~mdw / lisp / commitdiff