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     [ V{ } clone ] dip 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     [ up ] dip data-nth ; inline
  56
  57 : left-value ( n heap -- entry )
  58     [ left ] dip data-nth ; inline
  59
  60 : right-value ( n heap -- entry )
  61     [ right ] dip data-nth ; inline
  62
  63 : data-set-nth ( entry n heap -- )
  64     [ [ >>index drop ] 2keep ] dip
  65     data>> set-nth-unsafe ; inline
  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 [ data-nth ] dip ] 3keep
  86     tuck [ data-set-nth ] 2dip data-set-nth ; inline
  87
  88 GENERIC: heap-compare ( pair1 pair2 heap -- ? )
  89
  90 : (heap-compare) ( pair1 pair2 heap -- <=> )
  91     drop [ key>> ] compare ; inline
  92
  93 M: min-heap heap-compare (heap-compare) +gt+ eq? ;
  94
  95 M: max-heap heap-compare (heap-compare) +lt+ eq? ;
  96
  97 : heap-bounds-check? ( m heap -- ? )
  98     heap-size >= ; inline
  99
 100 : left-bounds-check? ( m heap -- ? )
 101     [ left ] dip heap-bounds-check? ; inline
 102
 103 : right-bounds-check? ( m heap -- ? )
 104     [ right ] dip heap-bounds-check? ; inline
 105
 106 : continue? ( m up[m] heap -- ? )
 107     [ data-nth swap ] keep [ data-nth ] keep
 108     heap-compare ; inline
 109
 110 DEFER: up-heap
 111
 112 : (up-heap) ( n heap -- )
 113     [ dup up ] dip
 114     3dup continue? [
 115         [ data-exchange ] 2keep up-heap
 116     ] [
 117         3drop
 118     ] if ;
 119
 120 : up-heap ( n heap -- )
 121     over 0 > [ (up-heap) ] [ 2drop ] if ;
 122
 123 : (child) ( m heap -- n )
 124     2dup right-value
 125     [ 2dup left-value ] dip
 126     rot heap-compare
 127     [ right ] [ left ] if ;
 128
 129 : child ( m heap -- n )
 130     2dup right-bounds-check?
 131     [ drop left ] [ (child) ] if ;
 132
 133 : swap-down ( m heap -- )
 134     [ child ] 2keep data-exchange ;
 135
 136 DEFER: down-heap
 137
 138 : (down-heap) ( m heap -- )
 139     [ child ] 2keep swapd
 140     3dup continue? [
 141         3drop
 142     ] [
 143         [ data-exchange ] 2keep down-heap
 144     ] if ;
 145
 146 : down-heap ( m heap -- )
 147     2dup left-bounds-check? [ 2drop ] [ (down-heap) ] if ;
 148
 149 PRIVATE>
 150
 151 M: heap heap-push* ( value key heap -- entry )
 152     [ <entry> dup ] keep [ data-push ] keep up-heap ;
 153
 154 : heap-push ( value key heap -- ) heap-push* drop ;
 155
 156 : heap-push-all ( assoc heap -- )
 157     [ swapd heap-push ] curry assoc-each ;
 158
 159 : >entry< ( entry -- key value )
 160     [ value>> ] [ key>> ] bi ;
 161
 162 M: heap heap-peek ( heap -- value key )
 163     data-first >entry< ;
 164
 165 ERROR: bad-heap-delete ;
 166
 167 M: bad-heap-delete summary
 168     drop "Invalid entry passed to heap-delete" ;
 169
 170 : entry>index ( entry heap -- n )
 171     over heap>> eq? [ bad-heap-delete ] unless
 172     index>> ;
 173
 174 M: heap heap-delete ( entry heap -- )
 175     [ entry>index ] keep
 176     2dup heap-size 1- = [
 177         nip data-pop*
 178     ] [
 179         [ nip data-pop ] 2keep
 180         [ data-set-nth ] 2keep
 181         down-heap
 182     ] if ;
 183
 184 M: heap heap-pop* ( heap -- )
 185     dup data-first swap heap-delete ;
 186
 187 M: heap heap-pop ( heap -- value key )
 188     dup data-first [ swap heap-delete ] keep >entry< ;
 189
 190 : heap-pop-all ( heap -- alist )
 191     [ dup heap-empty? not ]
 192     [ dup heap-pop swap 2array ]
 193     [ ] produce nip ;
 194
 195 : slurp-heap ( heap quot: ( elt -- ) -- )
 196     over heap-empty? [ 2drop ] [
 197         [ [ heap-pop drop ] dip call ] [ slurp-heap ] 2bi
 198     ] if ; inline recursive