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: accessors arrays assocs fry kernel kernel.private locals
   5 math math.order math.private sequences sequences.private summary
   6 vectors ;
   7 IN: heaps
   8
   9 GENERIC: heap-push* ( value key heap -- entry )
  10 GENERIC: heap-peek ( heap -- value key )
  11 GENERIC: heap-pop* ( heap -- )
  12 GENERIC: heap-pop ( heap -- value key )
  13 GENERIC: heap-delete ( entry heap -- )
  14 GENERIC: heap-empty? ( heap -- ? )
  15 GENERIC: heap-size ( heap -- n )
  16
  17 <PRIVATE
  18
  19 TUPLE: heap { data vector } ;
  20
  21 : <heap> ( class -- heap )
  22     V{ } clone swap boa ; inline
  23
  24 ERROR: not-a-heap object ;
  25
  26 : check-heap ( heap -- heap )
  27     dup heap? [ not-a-heap ] unless ; inline
  28
  29 TUPLE: entry value key heap index ;
  30
  31 : <entry> ( value key heap -- entry )
  32     f entry boa ; inline
  33
  34 PRIVATE>
  35
  36 TUPLE: min-heap < heap ;
  37
  38 : <min-heap> ( -- min-heap ) min-heap <heap> ;
  39
  40 TUPLE: max-heap < heap ;
  41
  42 : <max-heap> ( -- max-heap ) max-heap <heap> ;
  43
  44 M: heap heap-empty? ( heap -- ? )
  45     data>> empty? ; inline
  46
  47 M: heap heap-size ( heap -- n )
  48     data>> length ; inline
  49
  50 <PRIVATE
  51
  52 : left ( n -- m )
  53     { fixnum } declare 1 fixnum-shift-fast 1 fixnum+fast ; inline
  54
  55 : right ( n -- m )
  56     { fixnum } declare 1 fixnum-shift-fast 2 fixnum+fast ; inline
  57
  58 : up ( n -- m )
  59     { fixnum } declare 1 fixnum-fast 2/ ; inline
  60
  61 : data-nth ( n heap -- entry )
  62     data>> nth-unsafe { entry } declare ; inline
  63
  64 : data-first ( heap -- entry )
  65     0 swap data-nth ; inline
  66
  67 : data-set-nth ( entry n heap -- )
  68     [ [ >>index ] keep ] dip data>> set-nth-unsafe ; inline
  69
  70 : data-push ( entry heap -- n )
  71     [ heap-size [ >>index ] keep ]
  72     [ data>> [ set-nth ] 2keep drop ] bi ; inline
  73
  74 GENERIC: heap-compare ( entry1 entry2 heap -- ? )
  75
  76 M: min-heap heap-compare
  77     drop { entry entry } declare [ key>> ] bi@ after? ; inline
  78
  79 M: max-heap heap-compare
  80     drop { entry entry } declare [ key>> ] bi@ before? ; inline
  81
  82 : data-compare ( m n heap -- ? )
  83     [ '[ _ data-nth ] bi@ ] [ heap-compare ] bi ; inline
  84
  85 PRIVATE>
  86
  87 : >entry< ( entry -- value key )
  88     [ value>> ] [ key>> ] bi ; inline
  89
  90 M: heap heap-peek ( heap -- value key )
  91     data-first >entry< ;
  92
  93 <PRIVATE
  94
  95 :: sift-down ( heap from to -- )
  96     to heap data-nth :> tmp
  97
  98     to t [ over from > and ] [
  99         dup up
 100         dup heap data-nth
 101         dup tmp heap heap-compare [
 102             rot heap data-set-nth t
 103         ] [
 104             2drop f
 105         ] if
 106     ] while
 107
 108     tmp swap heap data-set-nth ; inline
 109
 110 PRIVATE>
 111
 112 M: heap heap-push*
 113     [ <entry> dup ] [ data-push ] [ 0 rot sift-down ] tri ;
 114
 115 : heap-push ( value key heap -- )
 116     heap-push* drop ;
 117
 118 : heap-push-all ( assoc heap -- )
 119     '[ swap _ heap-push ] assoc-each ;
 120
 121 <PRIVATE
 122
 123 :: sift-up ( heap n -- )
 124     heap heap-size  :> end
 125     n heap data-nth :> tmp
 126
 127     n dup left [ dup end < ] [
 128         dup 1 fixnum+fast
 129         dup end < [ 2dup heap data-compare ] [ f ] if
 130         [ nip ] [ drop ] if
 131         [ heap data-nth swap heap data-set-nth ]
 132         [ dup left ] bi
 133     ] while drop
 134
 135     tmp over heap data-set-nth
 136     heap n rot sift-down ; inline
 137
 138 PRIVATE>
 139
 140 M: heap heap-pop*
 141     dup data>> dup length 1 > [
 142         [ pop ] [ set-first ] bi 0 sift-up
 143     ] [
 144         pop* drop
 145     ] if ; inline
 146
 147 M: heap heap-pop
 148     [ data-first >entry< ] [ heap-pop* ] bi ;
 149
 150 : heap-pop-all ( heap -- alist )
 151     check-heap [ heap-size ] keep
 152     '[ _ heap-pop swap 2array ] replicate ;
 153
 154 : slurp-heap ( heap quot: ( value key -- ) -- )
 155     [ check-heap [ heap-size ] keep ] dip
 156     '[ _ heap-pop @ ] times ; inline
 157
 158 ERROR: bad-heap-delete ;
 159
 160 M: bad-heap-delete summary
 161     drop "Invalid entry passed to heap-delete" ;
 162
 163 <PRIVATE
 164
 165 : entry>index ( entry heap -- n )
 166     over heap>> eq? [ bad-heap-delete ] unless
 167     index>> { fixnum } declare ; inline
 168
 169 PRIVATE>
 170
 171 M: heap heap-delete
 172     [ entry>index ] keep
 173     2dup heap-size 1 - = [
 174         nip data>> pop*
 175     ] [
 176         [ nip data>> pop ]
 177         [ data-set-nth ]
 178         [ swap sift-up ] 2tri
 179     ] if ;
 180
 181 : >min-heap ( assoc -- min-heap )
 182     <min-heap> [ heap-push-all ] keep ;
 183
 184 : >max-heap ( assoc -- max-heap )
 185     <max-heap> [ heap-push-all ] keep ;