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 fry combinators ;
   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 )
  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? ( heap -- ? )
  39     data>> empty? ; inline
  40
  41 M: heap heap-size ( heap -- n )
  42     data>> length ; inline
  43
  44 <PRIVATE
  45
  46 : left ( n -- m ) 1 shift 1 + ; inline
  47
  48 : right ( n -- m ) 1 shift 2 + ; inline
  49
  50 : up ( n -- m ) 1 - 2/ ; inline
  51
  52 : data-nth ( n heap -- entry )
  53     data>> nth-unsafe ; inline
  54
  55 : left-value ( n heap -- entry )
  56     [ left ] dip data-nth ; inline
  57
  58 : right-value ( n heap -- entry )
  59     [ right ] dip data-nth ; inline
  60
  61 : data-set-nth ( entry n heap -- )
  62     [ [ >>index drop ] [ ] 2bi ] dip
  63     data>> set-nth-unsafe ; inline
  64
  65 : data-push ( entry heap -- n )
  66     dup heap-size [
  67         swap 2dup data>> ensure 2drop data-set-nth
  68     ] [
  69     ] bi ; inline
  70
  71 : data-first ( heap -- entry )
  72     data>> first ; inline
  73
  74 : data-exchange ( m n heap -- )
  75     [ '[ _ data-nth ] bi@ ]
  76     [ '[ _ data-set-nth ] bi@ ] 3bi ; inline
  77
  78 GENERIC: heap-compare ( entry1 entry2 heap -- ? )
  79
  80 : (heap-compare) ( entry1 entry2 heap -- <=> )
  81     drop [ key>> ] compare ; inline
  82
  83 M: min-heap heap-compare (heap-compare) +gt+ eq? ;
  84
  85 M: max-heap heap-compare (heap-compare) +lt+ eq? ;
  86
  87 : heap-bounds-check? ( m heap -- ? )
  88     heap-size >= ; inline
  89
  90 : left-bounds-check? ( m heap -- ? )
  91     [ left ] dip heap-bounds-check? ; inline
  92
  93 : right-bounds-check? ( m heap -- ? )
  94     [ right ] dip heap-bounds-check? ; inline
  95
  96 : continue? ( m n heap -- ? )
  97     [ data-nth nip ]
  98     [ nip data-nth ]
  99     [ 2nip ] 3tri heap-compare ; inline
 100
 101 DEFER: up-heap
 102
 103 : (up-heap) ( n heap -- )
 104     [ dup up ] dip
 105     3dup continue? [
 106         [ data-exchange ] [ up-heap ] 2bi
 107     ] [
 108         3drop
 109     ] if ; inline recursive
 110
 111 : up-heap ( n heap -- )
 112     over 0 > [ (up-heap) ] [ 2drop ] if ; inline recursive
 113
 114 : (child) ( m heap -- n )
 115     { [ drop ] [ left-value ] [ right-value ] [ nip ] } 2cleave
 116     heap-compare [ right ] [ left ] if ; inline
 117
 118 : child ( m heap -- n )
 119     2dup right-bounds-check?
 120     [ drop left ] [ (child) ] if ; inline
 121
 122 DEFER: down-heap
 123
 124 : (down-heap) ( m heap -- )
 125     [ drop ] [ child ] [ nip ] 2tri
 126     3dup continue? [
 127         3drop
 128     ] [
 129         [ data-exchange ] [ down-heap ] 2bi
 130     ] if ; inline recursive
 131
 132 : down-heap ( m heap -- )
 133     2dup left-bounds-check? [ 2drop ] [ (down-heap) ] if ; inline recursive
 134
 135 PRIVATE>
 136
 137 M: heap heap-push* ( value key heap -- entry )
 138     [ <entry> dup ] [ data-push ] [ ] tri up-heap ;
 139
 140 : heap-push ( value key heap -- ) heap-push* drop ;
 141
 142 : heap-push-all ( assoc heap -- )
 143     '[ swap _ heap-push ] assoc-each ;
 144
 145 : >entry< ( entry -- value key )
 146     [ value>> ] [ key>> ] bi ; inline
 147
 148 M: heap heap-peek ( heap -- value key )
 149     data-first >entry< ;
 150
 151 ERROR: bad-heap-delete ;
 152
 153 M: bad-heap-delete summary
 154     drop "Invalid entry passed to heap-delete" ;
 155
 156 : entry>index ( entry heap -- n )
 157     over heap>> eq? [ bad-heap-delete ] unless
 158     index>> ;
 159
 160 M: heap heap-delete ( entry heap -- )
 161     [ entry>index ] [ ] bi
 162     2dup heap-size 1 - = [
 163         nip data>> pop*
 164     ] [
 165         [ nip data>> pop ]
 166         [ data-set-nth ]
 167         [ ] 2tri
 168         down-heap
 169     ] if ;
 170
 171 M: heap heap-pop* ( heap -- )
 172     [ data-first ] keep heap-delete ;
 173
 174 M: heap heap-pop ( heap -- value key )
 175     [ data-first ] keep
 176     [ heap-delete ] [ drop ] 2bi >entry< ;
 177
 178 : heap-pop-all ( heap -- alist )
 179     [ dup heap-empty? not ]
 180     [ dup heap-pop swap 2array ]
 181     produce nip ;
 182
 183 : slurp-heap ( heap quot: ( elt -- ) -- )
 184     over heap-empty? [ 2drop ] [
 185         [ [ heap-pop drop ] dip call ] [ slurp-heap ] 2bi
 186     ] if ; inline recursive