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