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