mdw@git.distorted.org.uk Git - lisp/blob - heap.lisp

   1 ;;; -*-lisp-*-
   2 ;;;
   3 ;;; Heap data structure; useful for priority queues and suchlike
   4 ;;;
   5 ;;; (c) 2006 Straylight/Edgeware
   6 ;;;
   7
   8 ;;;----- Licensing notice ---------------------------------------------------
   9 ;;;
  10 ;;; This program is free software; you can redistribute it and/or modify
  11 ;;; it under the terms of the GNU General Public License as published by
  12 ;;; the Free Software Foundation; either version 2 of the License, or
  13 ;;; (at your option) any later version.
  14 ;;;
  15 ;;; This program is distributed in the hope that it will be useful,
  16 ;;; but WITHOUT ANY WARRANTY; without even the implied warranty of
  17 ;;; MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
  18 ;;; GNU General Public License for more details.
  19 ;;;
  20 ;;; You should have received a copy of the GNU General Public License
  21 ;;; along with this program; if not, write to the Free Software Foundation,
  22 ;;; Inc., 59 Temple Place - Suite 330, Boston, MA 02111-1307, USA.
  23
  24 (defpackage #:heap
  25   (:use #:common-lisp))
  26 (in-package #:heap)
  27
  28 ;;;--------------------------------------------------------------------------
  29 ;;; Useful indexing functions.
  30
  31 (declaim (inline parent left-child right-child))
  32 (deftype index () '(and unsigned-byte fixnum))
  33 (defun parent (i)
  34   (declare (type index i))
  35   (the index (floor (- i 1) 2)))
  36 (defun left-child (i)
  37   (declare (type index i))
  38   (the index (+ (* 2 i) 1)))
  39 (defun right-child (i)
  40   (declare (type index i))
  41   (the index (+ (* 2 i) 2)))
  42
  43 ;;;--------------------------------------------------------------------------
  44 ;;; Low-level heap operations.
  45
  46 (defun upheap (v key cmp n x)
  47   "Insert the element X in the highest place possible in the heap."
  48   (declare (type vector v)
  49            (type function key cmp)
  50            (type index n))
  51   (let ((i n) (xk (funcall key x)))
  52     (loop (when (zerop i) (return))
  53           (let* ((j (parent i))
  54                  (y (aref v j)))
  55             (when (funcall cmp (funcall key y) xk) (return))
  56             (setf (aref v i) y
  57                   i j)))
  58     (setf (aref v i) x)))
  59
  60 (defun downheap (v key cmp n x)
  61   "Insert the element X in the lowest place possible in the heap."
  62   (declare (type vector v)
  63            (type function key cmp)
  64            (type index n))
  65   (let ((i 0) (xk (funcall key x)))
  66     (loop (let ((l (left-child i))
  67                 (r (right-child i)))
  68             (when (>= l n) (return))
  69             (multiple-value-bind
  70                 (j y yk)
  71                 (let* ((y (aref v l))
  72                        (yk (funcall key y)))
  73                   (if (= r n)
  74                       (values l y yk)
  75                       (let* ((z (aref v r))
  76                              (zk (funcall key z)))
  77                         (if (funcall cmp yk zk)
  78                             (values l y yk)
  79                             (values r z zk)))))
  80               (when (funcall cmp xk yk)
  81                 (return))
  82               (setf (aref v i) y
  83                     i j))))
  84     (setf (aref v i) x)))
  85
  86 (defun check (v key cmp n)
  87   "Verify the heap invariant on the heap."
  88   (declare (type vector v)
  89            (type function key cmp)
  90            (type index n))
  91   (dotimes (i n)
  92     (let* ((item (aref v i))
  93            (item-key (funcall key item))
  94            (l (left-child i))
  95            (r (right-child i)))
  96       (when (< l n)
  97         (let ((left-item (aref v l)))
  98           (assert (funcall cmp item-key (funcall key left-item))))
  99         (when (< r n)
 100           (let ((right-item (aref v r)))
 101             (assert (funcall cmp item-key (funcall key right-item)))))))))
 102
 103 ;;;--------------------------------------------------------------------------
 104 ;;; High-level heap things
 105
 106 (export '(heap heapp))
 107 (defstruct (heap (:predicate heapp) (:constructor %make-heap))
 108   "Data structure for a heap."
 109   (v (make-array 16) :type vector)
 110   (n 0 :type index)
 111   (key #'identity :type function :read-only t)
 112   (compare #'<= :type function :read-only t))
 113
 114 (export 'make-heap)
 115 (defun make-heap
 116     (&key (compare #'<=) (key #'identity)
 117           (type 't) (init-size 16) (contents nil contentsp))
 118   "Return a new heap.
 119
 120    COMPARE is a partial-order predicate: (COMPARE X Y) should return true if
 121    X <= Y in some order.
 122
 123    The TYPE is the element type of the heap.
 124
 125    INIT-SIZE is the initial allocation for the heap; the heap will grow
 126    automatically if necessary, so this isn't a big deal.  This is only a
 127    hint; make-heap may ignore it completely.
 128
 129    KEY is a function to extract the key from an element.  The default is to
 130    use the item unmolested.
 131
 132    CONTENTS is the initial contents of the heap.  If omitted, the heap is
 133    initially empty."
 134   (let ((n (if contentsp (length contents) 0)))
 135     (loop while (< init-size n)
 136           do (setf init-size (ash init-size 1)))
 137     (let ((v (make-array init-size :element-type type)))
 138       (when contentsp
 139         (reduce (lambda (i item)
 140                   (upheap v key compare i item)
 141                   (1+ i))
 142                 contents
 143                 :initial-value 0))
 144       (%make-heap :compare compare :key key :n n :v v))))
 145
 146 (export 'heap-count)
 147 (defun heap-count (heap)
 148   "Return the number of elements in HEAP."
 149   (declare (type heap heap))
 150   (heap-n heap))
 151
 152 (export 'heap-empty-p)
 153 (defun heap-empty-p (heap)
 154   "True if HEAP is empty."
 155   (declare (type heap heap))
 156   (zerop (heap-count heap)))
 157
 158 (export 'heap-insert)
 159 (defun heap-insert (heap item)
 160   "Insert ITEM into the HEAP."
 161   (declare (type heap heap))
 162   (let* ((v (heap-v heap))
 163          (n (heap-n heap))
 164          (sz (array-dimension v 0)))
 165     (when (= n sz)
 166       (setf v (adjust-array v (* 2 n))
 167             (heap-v heap) v))
 168     (upheap v (heap-key heap) (heap-compare heap) n item)
 169     (setf (heap-n heap) (1+ n))))
 170
 171 (export 'heap-head)
 172 (defun heap-head (heap)
 173   "Peep at the head item on HEAP."
 174   (declare (type heap heap))
 175   (assert (not (heap-empty-p heap)))
 176   (aref (heap-v heap) 0))
 177
 178 (export 'heap-remove)
 179 (defun heap-remove (heap)
 180   "Remove the head item from HEAP and return it."
 181   (declare (type heap heap))
 182   (assert (not (heap-empty-p heap)))
 183   (let ((v (heap-v heap))
 184         (n (1- (heap-n heap))))
 185     (prog1 (aref v 0)
 186       (setf (heap-n heap) n)
 187       (downheap v (heap-key heap) (heap-compare heap) n (aref v n)))))
 188
 189 (export 'heap-sort)
 190 (defun heap-sort (items compare &key (key #'identity))
 191   "Return the ITEMS, least-first, as sorted by the ordering COMPARE."
 192   (let ((heap (make-heap :compare compare :contents items :key key)))
 193     (loop repeat (heap-n heap)
 194           collect (heap-remove heap))))
 195
 196 ;;;----- That's all, folks --------------------------------------------------