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