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