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? data>> empty? ; inline
  39
  40 M: heap heap-size data>> length ; inline
  41
  42 <PRIVATE
  43
  44 : left ( n -- m )
  45     { fixnum } declare 1 fixnum-shift-fast 1 fixnum+fast ; inline
  46
  47 : right ( n -- m )
  48     { fixnum } declare 1 fixnum-shift-fast 2 fixnum+fast ; inline
  49
  50 : up ( n -- m )
  51     { fixnum } declare 1 fixnum-fast 2/ ; inline
  52
  53 : data-nth ( n data -- entry )
  54     nth-unsafe { entry } declare ; inline
  55
  56 : data-set-nth ( entry n data -- )
  57     [ [ >>index ] keep ] dip set-nth-unsafe ; inline
  58
  59 : data-push ( entry data -- n )
  60     [ length [ >>index ] keep ]
  61     [ [ set-nth ] keepd ] bi ; inline
  62
  63 GENERIC: heap-compare ( entry1 entry2 heap -- ? )
  64
  65 M: min-heap heap-compare
  66     drop { entry entry } declare [ key>> ] bi@ after? ; inline
  67
  68 M: max-heap heap-compare
  69     drop { entry entry } declare [ key>> ] bi@ before? ; inline
  70
  71 PRIVATE>
  72
  73 : >entry< ( entry -- value key )
  74     [ value>> ] [ key>> ] bi ; inline
  75
  76 M: heap heap-peek
  77     data>> first >entry< ;
  78
  79 <PRIVATE
  80
  81 ! names and optimizations inspired by cpython/Lib/heapq.py,
  82 ! refer to it from in depth explanations of the optimizations.
  83
  84 ! called bubble-up in the literature... but we keep cpython's name.
  85 :: sift-down ( heap from to -- )
  86     heap data>>      :> data
  87     to data data-nth :> tmp
  88
  89     to t [ over from > and ] [
  90         dup up
  91         dup data data-nth
  92         dup tmp heap heap-compare [
  93             rot data data-set-nth t
  94         ] [
  95             2drop f
  96         ] if
  97     ] while
  98
  99     tmp swap data data-set-nth ; inline
 100
 101 PRIVATE>
 102
 103 M: heap heap-push*
 104     [ <entry> dup ] [ data>> data-push ] [ 0 rot sift-down ] tri ;
 105
 106 : heap-push ( value key heap -- )
 107     heap-push* drop ;
 108
 109 : heap-push-all ( assoc heap -- )
 110     '[ swap _ heap-push ] assoc-each ;
 111
 112 <PRIVATE
 113
 114 ! called bubble-down in the literature... but we keep cpython's name.
 115 ! A quote from cpython's implementation:
 116 ! > We *could* break out of the loop as soon as we find a pos where newitem <=
 117 ! > both its children, but turns out that's not a good idea [...]
 118 ! Indeed the code is 33% slower if we remove this optimization.
 119 :: sift-up ( heap n -- )
 120     heap data>>     :> data
 121     data length     :> end
 122     n data data-nth :> tmp
 123
 124     n dup left [ dup end < ] [
 125         dup 1 fixnum+fast
 126         dup end < [
 127             2dup [ data data-nth ] bi@ heap heap-compare
 128         ] [ f ] if
 129         [ nip ] [ drop ] if
 130         [ data data-nth swap data data-set-nth ]
 131         [ dup left ] bi
 132     ] while drop
 133
 134     tmp over data data-set-nth
 135     heap n rot sift-down ; inline
 136
 137 PRIVATE>
 138
 139 M: heap heap-pop*
 140     dup data>> f over first index<< [ pop ] keep
 141     [ 2drop ] [ set-first 0 sift-up ] if-empty ;
 142
 143 : heap-pop ( heap -- value key )
 144     [ heap-peek ] [ heap-pop* ] bi ;
 145
 146 : slurp-heap ( ... heap quot: ( ... value key -- ... ) -- ... )
 147     [ drop '[ _ heap-empty? ] ]
 148     [ '[ _ heap-pop @ ] until ] 2bi ; inline
 149
 150 : heap-pop-all ( heap -- alist )
 151     [ heap-size <vector> ] keep
 152     [ swap 2array suffix! ] slurp-heap { } like ;
 153
 154 ERROR: bad-heap-delete ;
 155
 156 M: bad-heap-delete summary
 157     drop "Invalid entry passed to heap-delete" ;
 158
 159 <PRIVATE
 160
 161 : entry>index ( entry heap -- n )
 162     over heap>> eq? [ bad-heap-delete ] unless
 163     index>> dup [ bad-heap-delete ] unless
 164     { fixnum } declare ; inline
 165
 166 PRIVATE>
 167
 168 M:: heap heap-delete ( entry heap -- )
 169     entry heap entry>index :> n
 170     heap data>>            :> data
 171     data pop               :> nth-entry
 172     f entry index<<
 173     n data length = [
 174         nth-entry n data data-set-nth
 175         n 0 = [ t ] [ nth-entry n up data data-nth heap heap-compare ] if
 176         [ heap n sift-up ] [ heap 0 n sift-down ] if
 177     ] unless ;
 178
 179 : >min-heap ( assoc -- min-heap )
 180     dup assoc-size <vector> min-heap boa
 181     [ heap-push-all ] keep ;
 182
 183 : >max-heap ( assoc -- max-heap )
 184     dup assoc-size <vector> max-heap boa
 185     [ heap-push-all ] keep ;