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: kernel math sequences arrays assocs sequences.private
   5 growable accessors math.order summary vectors ;
   6 IN: heaps
   7
   8 GENERIC: heap-push* ( value key heap -- entry )
   9 GENERIC: heap-peek ( heap -- value key )
  10 GENERIC: heap-pop* ( heap -- )
  11 GENERIC: heap-pop ( heap -- value key )
  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 ] dip boa ; inline
  22
  23 TUPLE: entry value key heap index ;
  24
  25 : <entry> ( value key heap -- entry ) f entry boa ; inline
  26
  27 PRIVATE>
  28
  29 TUPLE: min-heap < heap ;
  30
  31 : <min-heap> ( -- min-heap ) min-heap <heap> ;
  32
  33 TUPLE: max-heap < heap ;
  34
  35 : <max-heap> ( -- max-heap ) max-heap <heap> ;
  36
  37 M: heap heap-empty? ( heap -- ? )
  38     data>> empty? ;
  39
  40 M: heap heap-size ( heap -- n )
  41     data>> length ;
  42
  43 <PRIVATE
  44
  45 : left ( n -- m ) 1 shift 1 + ; inline
  46
  47 : right ( n -- m ) 1 shift 2 + ; inline
  48
  49 : up ( n -- m ) 1 - 2/ ; inline
  50
  51 : data-nth ( n heap -- entry )
  52     data>> nth-unsafe ; inline
  53
  54 : left-value ( n heap -- entry )
  55     [ left ] dip data-nth ; inline
  56
  57 : right-value ( n heap -- entry )
  58     [ right ] dip data-nth ; inline
  59
  60 : data-set-nth ( entry n heap -- )
  61     [ [ >>index drop ] 2keep ] dip
  62     data>> set-nth-unsafe ; inline
  63
  64 : data-push ( entry heap -- n )
  65     dup heap-size [
  66         swap 2dup data>> ensure 2drop data-set-nth
  67     ] keep ; inline
  68
  69 : data-pop ( heap -- entry )
  70     data>> pop ; inline
  71
  72 : data-pop* ( heap -- )
  73     data>> pop* ; inline
  74
  75 : data-first ( heap -- entry )
  76     data>> first ; inline
  77
  78 : data-exchange ( m n heap -- )
  79     [ [ data-nth ] curry bi@ ]
  80     [ [ data-set-nth ] curry bi@ ] 3bi ; inline
  81
  82 GENERIC: heap-compare ( pair1 pair2 heap -- ? )
  83
  84 : (heap-compare) ( pair1 pair2 heap -- <=> )
  85     drop [ key>> ] compare ; inline
  86
  87 M: min-heap heap-compare (heap-compare) +gt+ eq? ;
  88
  89 M: max-heap heap-compare (heap-compare) +lt+ eq? ;
  90
  91 : heap-bounds-check? ( m heap -- ? )
  92     heap-size >= ; inline
  93
  94 : left-bounds-check? ( m heap -- ? )
  95     [ left ] dip heap-bounds-check? ; inline
  96
  97 : right-bounds-check? ( m heap -- ? )
  98     [ right ] dip heap-bounds-check? ; inline
  99
 100 : continue? ( m up[m] heap -- ? )
 101     [ data-nth swap ] keep [ data-nth ] keep
 102     heap-compare ; inline
 103
 104 DEFER: up-heap
 105
 106 : (up-heap) ( n heap -- )
 107     [ dup up ] dip
 108     3dup continue? [
 109         [ data-exchange ] 2keep up-heap
 110     ] [
 111         3drop
 112     ] if ; inline recursive
 113
 114 : up-heap ( n heap -- )
 115     over 0 > [ (up-heap) ] [ 2drop ] if ; inline recursive
 116
 117 : (child) ( m heap -- n )
 118     2dup right-value
 119     [ 2dup left-value ] dip
 120     rot heap-compare
 121     [ right ] [ left ] if ;
 122
 123 : child ( m heap -- n )
 124     2dup right-bounds-check?
 125     [ drop left ] [ (child) ] if ;
 126
 127 DEFER: down-heap
 128
 129 : (down-heap) ( m heap -- )
 130     [ child ] 2keep swapd
 131     3dup continue? [
 132         3drop
 133     ] [
 134         [ data-exchange ] 2keep down-heap
 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 ] keep [ data-push ] keep up-heap ;
 144
 145 : heap-push ( value key heap -- ) heap-push* drop ;
 146
 147 : heap-push-all ( assoc heap -- )
 148     [ swapd heap-push ] curry assoc-each ;
 149
 150 : >entry< ( entry -- key value )
 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>> ;
 164
 165 M: heap heap-delete ( entry heap -- )
 166     [ entry>index ] keep
 167     2dup heap-size 1 - = [
 168         nip data-pop*
 169     ] [
 170         [ nip data-pop ] 2keep
 171         [ data-set-nth ] 2keep
 172         down-heap
 173     ] if ;
 174
 175 M: heap heap-pop* ( heap -- )
 176     dup data-first swap heap-delete ;
 177
 178 M: heap heap-pop ( heap -- value key )
 179     dup data-first [ swap heap-delete ] keep >entry< ;
 180
 181 : heap-pop-all ( heap -- alist )
 182     [ dup heap-empty? not ]
 183     [ dup heap-pop swap 2array ]
 184     produce nip ;
 185
 186 : heap-values ( heap -- alist )
 187     data>> [ value>> ] { } map-as ;
 188
 189 : slurp-heap ( heap quot: ( elt -- ) -- )
 190     over heap-empty? [ 2drop ] [
 191         [ [ heap-pop drop ] dip call ] [ slurp-heap ] 2bi
 192     ] if ; inline recursive