basis/heaps/heaps.factor

   1 ! Copyright (C) 2007, 2008 Ryan Murphy, Doug Coleman,
   2 ! Slava Pestov.
   3 ! See http://factorcode.org/license.txt for BSD license.
   4 USING: kernel math sequences arrays assocs sequences.private
   5 growable accessors math.order summary ;
   6 IN: heaps
   7
   8 GENERIC: heap-push* ( value key heap -- entry )
   9 GENERIC: heap-peek ( heap -- value key )
  10 GENERIC: heap-pop* ( heap -- )
  11 GENERIC: heap-pop ( heap -- value key )
  12 GENERIC: heap-delete ( entry heap -- )
  13 GENERIC: heap-empty? ( heap -- ? )
  14 GENERIC: heap-size ( heap -- n )
  15
  16 <PRIVATE
  17
  18 TUPLE: heap data ;
  19
  20 : <heap> ( class -- heap )
  21     >r V{ } clone r> boa ; inline
  22
  23 TUPLE: entry value key heap index ;
  24
  25 : <entry> ( value key heap -- entry ) f entry boa ;
  26
  27 PRIVATE>
  28
  29 TUPLE: min-heap < heap ;
  30
  31 : <min-heap> ( -- min-heap ) min-heap <heap> ;
  32
  33 TUPLE: max-heap < heap ;
  34
  35 : <max-heap> ( -- max-heap ) max-heap <heap> ;
  36
  37 M: heap heap-empty? ( heap -- ? )
  38     data>> empty? ;
  39
  40 M: heap heap-size ( heap -- n )
  41     data>> length ;
  42
  43 <PRIVATE
  44
  45 : left ( n -- m ) 1 shift 1 + ; inline
  46
  47 : right ( n -- m ) 1 shift 2 + ; inline
  48
  49 : up ( n -- m ) 1- 2/ ; inline
  50
  51 : data-nth ( n heap -- entry )
  52     data>> nth-unsafe ; inline
  53
  54 : up-value ( n heap -- entry )
  55     >r up r> data-nth ; inline
  56
  57 : left-value ( n heap -- entry )
  58     >r left r> data-nth ; inline
  59
  60 : right-value ( n heap -- entry )
  61     >r right r> data-nth ; inline
  62
  63 : data-set-nth ( entry n heap -- )
  64     >r [ >>index drop ] 2keep r>
  65     data>> set-nth-unsafe ;
  66
  67 : data-push ( entry heap -- n )
  68     dup heap-size [
  69         swap 2dup data>> ensure 2drop data-set-nth
  70     ] keep ; inline
  71
  72 : data-pop ( heap -- entry )
  73     data>> pop ; inline
  74
  75 : data-pop* ( heap -- )
  76     data>> pop* ; inline
  77
  78 : data-peek ( heap -- entry )
  79     data>> peek ; inline
  80
  81 : data-first ( heap -- entry )
  82     data>> first ; inline
  83
  84 : data-exchange ( m n heap -- )
  85     [ tuck data-nth >r data-nth r> ] 3keep
  86     tuck >r >r data-set-nth r> r> data-set-nth ; inline
  87
  88 GENERIC: heap-compare ( pair1 pair2 heap -- ? )
  89
  90 : (heap-compare) drop [ key>> ] compare ; inline
  91
  92 M: min-heap heap-compare (heap-compare) +gt+ eq? ;
  93
  94 M: max-heap heap-compare (heap-compare) +lt+ eq? ;
  95
  96 : heap-bounds-check? ( m heap -- ? )
  97     heap-size >= ; inline
  98
  99 : left-bounds-check? ( m heap -- ? )
 100     >r left r> heap-bounds-check? ; inline
 101
 102 : right-bounds-check? ( m heap -- ? )
 103     >r right r> heap-bounds-check? ; inline
 104
 105 : continue? ( m up[m] heap -- ? )
 106     [ data-nth swap ] keep [ data-nth ] keep
 107     heap-compare ; inline
 108
 109 DEFER: up-heap
 110
 111 : (up-heap) ( n heap -- )
 112     >r dup up r>
 113     3dup continue? [
 114         [ data-exchange ] 2keep up-heap
 115     ] [
 116         3drop
 117     ] if ;
 118
 119 : up-heap ( n heap -- )
 120     over 0 > [ (up-heap) ] [ 2drop ] if ;
 121
 122 : (child) ( m heap -- n )
 123     2dup right-value
 124     >r 2dup left-value r>
 125     rot heap-compare
 126     [ right ] [ left ] if ;
 127
 128 : child ( m heap -- n )
 129     2dup right-bounds-check?
 130     [ drop left ] [ (child) ] if ;
 131
 132 : swap-down ( m heap -- )
 133     [ child ] 2keep data-exchange ;
 134
 135 DEFER: down-heap
 136
 137 : (down-heap) ( m heap -- )
 138     [ child ] 2keep swapd
 139     3dup continue? [
 140         3drop
 141     ] [
 142         [ data-exchange ] 2keep down-heap
 143     ] if ;
 144
 145 : down-heap ( m heap -- )
 146     2dup left-bounds-check? [ 2drop ] [ (down-heap) ] if ;
 147
 148 PRIVATE>
 149
 150 M: heap heap-push* ( value key heap -- entry )
 151     [ <entry> dup ] keep [ data-push ] keep up-heap ;
 152
 153 : heap-push ( value key heap -- ) heap-push* drop ;
 154
 155 : heap-push-all ( assoc heap -- )
 156     [ swapd heap-push ] curry assoc-each ;
 157
 158 : >entry< ( entry -- key value )
 159     [ value>> ] [ key>> ] bi ;
 160
 161 M: heap heap-peek ( heap -- value key )
 162     data-first >entry< ;
 163
 164 ERROR: bad-heap-delete ;
 165
 166 M: bad-heap-delete summary
 167     drop "Invalid entry passed to heap-delete" ;
 168
 169 : entry>index ( entry heap -- n )
 170     over heap>> eq? [ bad-heap-delete ] unless
 171     index>> ;
 172
 173 M: heap heap-delete ( entry heap -- )
 174     [ entry>index ] keep
 175     2dup heap-size 1- = [
 176         nip data-pop*
 177     ] [
 178         [ nip data-pop ] 2keep
 179         [ data-set-nth ] 2keep
 180         down-heap
 181     ] if ;
 182
 183 M: heap heap-pop* ( heap -- )
 184     dup data-first swap heap-delete ;
 185
 186 M: heap heap-pop ( heap -- value key )
 187     dup data-first [ swap heap-delete ] keep >entry< ;
 188
 189 : heap-pop-all ( heap -- alist )
 190     [ dup heap-empty? not ]
 191     [ dup heap-pop swap 2array ]
 192     [ ] produce nip ;
 193
 194 : slurp-heap ( heap quot: ( elt -- ) -- )
 195     over heap-empty? [ 2drop ] [
 196         [ [ heap-pop drop ] dip call ] [ slurp-heap ] 2bi
 197     ] if ; inline recursive