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 combinators fry kernel
   5 kernel.private locals math math.order math.private sequences
   6 sequences.private summary vectors ;
   7 IN: heaps
   8
   9 ! names and optimizations copied from pypy's heapq.py,
  10 ! refer to it from in depth explanations of the optimizations.
  11
  12 GENERIC: heap-push* ( value key heap -- entry )
  13 GENERIC: heap-peek ( heap -- value key )
  14 GENERIC: heap-pop* ( heap -- )
  15 GENERIC: heap-delete ( entry heap -- )
  16 GENERIC: heap-empty? ( heap -- ? )
  17 GENERIC: heap-size ( heap -- n )
  18
  19 <PRIVATE
  20
  21 TUPLE: heap { data vector } ;
  22
  23 : <heap> ( class -- heap )
  24     V{ } clone swap boa ; inline
  25
  26 TUPLE: entry value key heap index ;
  27
  28 : <entry> ( value key heap -- entry )
  29     f entry boa ; inline
  30
  31 PRIVATE>
  32
  33 TUPLE: min-heap < heap ;
  34
  35 : <min-heap> ( -- min-heap ) min-heap <heap> ;
  36
  37 TUPLE: max-heap < heap ;
  38
  39 : <max-heap> ( -- max-heap ) max-heap <heap> ;
  40
  41 M: heap heap-empty? ( heap -- ? )
  42     data>> empty? ; inline
  43
  44 M: heap heap-size ( heap -- n )
  45     data>> length ; inline
  46
  47 <PRIVATE
  48
  49 : left ( n -- m )
  50     { fixnum } declare 1 fixnum-shift-fast 1 fixnum+fast ; inline
  51
  52 : right ( n -- m )
  53     { fixnum } declare 1 fixnum-shift-fast 2 fixnum+fast ; inline
  54
  55 : up ( n -- m )
  56     { fixnum } declare 1 fixnum-fast 2/ ; inline
  57
  58 : data-nth ( n data -- entry )
  59     nth-unsafe { entry } declare ; inline
  60
  61 : data-set-nth ( entry n data -- )
  62     [ [ >>index ] keep ] dip set-nth-unsafe ; inline
  63
  64 : data-push ( entry data -- n )
  65     [ length [ >>index ] keep ]
  66     [ [ set-nth ] keepd ] bi ; inline
  67
  68 GENERIC: heap-compare ( entry1 entry2 heap -- ? )
  69
  70 M: min-heap heap-compare
  71     drop { entry entry } declare [ key>> ] bi@ after? ; inline
  72
  73 M: max-heap heap-compare
  74     drop { entry entry } declare [ key>> ] bi@ before? ; inline
  75
  76 : (heapdata-compare) ( m n data heap -- ? )
  77     [ '[ _ data-nth ] bi@ ] [ heap-compare ] bi* ; inline
  78
  79 : heapdata-compare ( m n heap -- ? )
  80     [ data>> ] keep (heapdata-compare) ; inline
  81
  82 PRIVATE>
  83
  84 : >entry< ( entry -- value key )
  85     [ value>> ] [ key>> ] bi ; inline
  86
  87 M: heap heap-peek ( heap -- value key )
  88     data>> first >entry< ;
  89
  90 <PRIVATE
  91
  92 ! called bubble-up in the litterature... but we keep pypy's name.
  93 :: sift-down ( heap from to -- )
  94     heap data>>      :> data
  95     to data data-nth :> tmp
  96
  97     to t [ over from > and ] [
  98         dup up
  99         dup data data-nth
 100         dup tmp heap heap-compare [
 101             rot data data-set-nth t
 102         ] [
 103             2drop f
 104         ] if
 105     ] while
 106
 107     tmp swap data data-set-nth ; inline
 108
 109 PRIVATE>
 110
 111 M: heap heap-push*
 112     [ <entry> dup ] [ data>> data-push ] [ 0 rot sift-down ] tri ;
 113
 114 : heap-push ( value key heap -- )
 115     heap-push* drop ;
 116
 117 : heap-push-all ( assoc heap -- )
 118     '[ swap _ heap-push ] assoc-each ;
 119
 120 <PRIVATE
 121
 122 ! called bubble-down in the litterature... but we keep pypy's name.
 123 ! A quote from pypy's implementation:
 124 ! > We *could* break out of the loop as soon as we find a pos where newitem <=
 125 ! > both its children, but turns out that's not a good idea [...]
 126 ! Indeed the code is 33% slower if we remove this optmization.
 127 :: sift-up ( heap n -- )
 128     heap data>>     :> data
 129     data length     :> end
 130     n data data-nth :> tmp
 131
 132     n dup left [ dup end < ] [
 133         dup 1 fixnum+fast
 134         dup end < [
 135             2dup data heap (heapdata-compare)
 136         ] [ f ] if
 137         [ nip ] [ drop ] if
 138         [ data data-nth swap data data-set-nth ]
 139         [ dup left ] bi
 140     ] while drop
 141
 142     tmp over data data-set-nth
 143     heap n rot sift-down ; inline
 144
 145 : (heap-pop*) ( heap data -- )
 146     [ first f >>index drop ] [ pop ] [ set-first ] tri 0 sift-up ; inline
 147
 148 PRIVATE>
 149
 150 M: heap heap-pop*
 151     dup data>> dup length 1 > [
 152         (heap-pop*)
 153     ] [
 154         pop f >>index 2drop
 155     ] if ; inline
 156
 157 : heap-pop ( heap -- value key )
 158     [ heap-peek ] [ heap-pop* ] bi ;
 159
 160 : slurp-heap ( ... heap quot: ( ... value key -- ... ) -- ... )
 161     [ drop '[ _ heap-empty? ] ]
 162     [ '[ _ heap-pop @ ] until ] 2bi ; inline
 163
 164 : heap-pop-all ( heap -- alist )
 165     [ heap-size <vector> ] keep
 166     [ swap 2array suffix! ] slurp-heap { } like ;
 167
 168 ERROR: bad-heap-delete ;
 169
 170 M: bad-heap-delete summary
 171     drop "Invalid entry passed to heap-delete" ;
 172
 173 <PRIVATE
 174
 175 : entry>index ( entry heap -- n )
 176     over heap>> eq? [ bad-heap-delete ] unless
 177     index>> dup [ bad-heap-delete ] unless
 178     { fixnum } declare ; inline
 179
 180 PRIVATE>
 181
 182 : ((heap-delete)) ( n heap -- )
 183     2dup [ dup up ] dip heapdata-compare
 184     [ swap sift-up ] [ 0 rot sift-down ] if ;
 185
 186 : (heap-delete) ( n heap -- )
 187     [ nip data>> pop ]
 188     [ data>> data-set-nth ]
 189     [ ((heap-delete)) ] 2tri ;
 190
 191 M: heap heap-delete
 192     [ entry>index ] [ f rot index<< ] 2bi
 193     {
 194         { [ 2dup heap-size 1 - = ] [ nip data>> pop* ] }
 195         { [ over zero? ] [ nip dup data>> (heap-pop*) ] }
 196         [ (heap-delete) ]
 197     } cond ;
 198
 199 : >min-heap ( assoc -- min-heap )
 200     dup assoc-size <vector> min-heap boa
 201     [ heap-push-all ] keep ;
 202
 203 : >max-heap ( assoc -- max-heap )
 204     dup assoc-size <vector> max-heap boa
 205     [ heap-push-all ] keep ;