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