1 ! Copyright (C) 2007, 2008 Ryan Murphy, Doug Coleman,
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
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 )
19 TUPLE: heap { data vector } ;
21 : <heap> ( class -- heap )
22 V{ } clone swap boa ; inline
24 ERROR: not-a-heap object ;
26 : check-heap ( heap -- heap )
27 dup heap? [ not-a-heap ] unless ; inline
29 TUPLE: entry value key heap index ;
31 : <entry> ( value key heap -- entry )
36 TUPLE: min-heap < heap ;
38 : <min-heap> ( -- min-heap ) min-heap <heap> ;
40 TUPLE: max-heap < heap ;
42 : <max-heap> ( -- max-heap ) max-heap <heap> ;
44 M: heap heap-empty? ( heap -- ? )
45 data>> empty? ; inline
47 M: heap heap-size ( heap -- n )
48 data>> length ; inline
53 { fixnum } declare 1 fixnum-shift-fast 1 fixnum+fast ; inline
56 { fixnum } declare 1 fixnum-shift-fast 2 fixnum+fast ; inline
59 { fixnum } declare 1 fixnum-fast 2/ ; inline
61 : data-nth ( n heap -- entry )
62 data>> nth-unsafe { entry } declare ; inline
64 : data-first ( heap -- entry )
67 : data-set-nth ( entry n heap -- )
68 [ [ >>index ] keep ] dip data>> set-nth-unsafe ; inline
70 : data-push ( entry heap -- n )
71 [ heap-size [ >>index ] keep ]
72 [ data>> [ set-nth ] 2keep drop ] bi ; inline
74 GENERIC: heap-compare ( entry1 entry2 heap -- ? )
76 M: min-heap heap-compare
77 drop { entry entry } declare [ key>> ] bi@ after? ; inline
79 M: max-heap heap-compare
80 drop { entry entry } declare [ key>> ] bi@ before? ; inline
82 : data-compare ( m n heap -- ? )
83 [ '[ _ data-nth ] bi@ ] [ heap-compare ] bi ; inline
87 : >entry< ( entry -- value key )
88 [ value>> ] [ key>> ] bi ; inline
90 M: heap heap-peek ( heap -- value key )
95 :: sift-down ( heap from to -- )
96 to heap data-nth :> tmp
98 to t [ over from > and ] [
101 dup tmp heap heap-compare [
102 rot heap data-set-nth t
108 tmp swap heap data-set-nth ; inline
113 [ <entry> dup ] [ data-push ] [ 0 rot sift-down ] tri ;
115 : heap-push ( value key heap -- )
118 : heap-push-all ( assoc heap -- )
119 '[ swap _ heap-push ] assoc-each ;
123 :: sift-up ( heap n -- )
124 heap heap-size :> end
125 n heap data-nth :> tmp
127 n dup left [ dup end < ] [
129 dup end < [ 2dup heap data-compare ] [ f ] if
131 [ heap data-nth swap heap data-set-nth ]
135 tmp over heap data-set-nth
136 heap n rot sift-down ; inline
141 dup data>> dup length 1 > [
142 [ pop ] [ set-first ] bi 0 sift-up
148 [ data-first >entry< ] [ heap-pop* ] bi ;
150 : slurp-heap ( ... heap quot: ( ... value key -- ... ) -- ... )
152 [ drop '[ _ heap-empty? ] ]
153 [ '[ _ heap-pop @ ] until ] 2bi ; inline
155 : heap-pop-all ( heap -- alist )
156 [ heap-size <vector> ] keep
157 [ swap 2array suffix! ] slurp-heap { } like ;
159 ERROR: bad-heap-delete ;
161 M: bad-heap-delete summary
162 drop "Invalid entry passed to heap-delete" ;
166 : entry>index ( entry heap -- n )
167 over heap>> eq? [ bad-heap-delete ] unless
168 index>> { fixnum } declare ; inline
174 2dup heap-size 1 - = [
179 [ swap sift-up ] 2tri
182 : >min-heap ( assoc -- min-heap )
183 dup assoc-size <vector> min-heap boa
184 [ heap-push-all ] keep ;
186 : >max-heap ( assoc -- max-heap )
187 dup assoc-size <vector> max-heap boa
188 [ heap-push-all ] keep ;