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

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