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-delete ( entry heap -- )
13 GENERIC: heap-empty? ( heap -- ? )
14 GENERIC: heap-size ( heap -- n )
18 TUPLE: heap { data vector } ;
20 : <heap> ( class -- heap )
21 V{ } clone swap boa ; inline
23 TUPLE: entry value key heap index ;
25 : <entry> ( value key heap -- entry )
30 TUPLE: min-heap < heap ;
32 : <min-heap> ( -- min-heap ) min-heap <heap> ;
34 TUPLE: max-heap < heap ;
36 : <max-heap> ( -- max-heap ) max-heap <heap> ;
38 M: heap heap-empty? ( heap -- ? )
39 data>> empty? ; inline
41 M: heap heap-size ( heap -- n )
42 data>> length ; inline
47 { fixnum } declare 1 fixnum-shift-fast 1 fixnum+fast ; inline
50 { fixnum } declare 1 fixnum-shift-fast 2 fixnum+fast ; inline
53 { fixnum } declare 1 fixnum-fast 2/ ; inline
55 : data-nth ( n data -- entry )
56 nth-unsafe { entry } declare ; inline
58 : data-set-nth ( entry n data -- )
59 [ [ >>index ] keep ] dip set-nth-unsafe ; inline
61 : data-push ( entry data -- n )
62 [ length [ >>index ] keep ]
63 [ [ set-nth ] keepd ] bi ; inline
65 GENERIC: heap-compare ( entry1 entry2 heap -- ? )
67 M: min-heap heap-compare
68 drop { entry entry } declare [ key>> ] bi@ after? ; inline
70 M: max-heap heap-compare
71 drop { entry entry } declare [ key>> ] bi@ before? ; inline
75 : >entry< ( entry -- value key )
76 [ value>> ] [ key>> ] bi ; inline
78 M: heap heap-peek ( heap -- value key )
79 data>> first >entry< ;
83 ! names and optimizations inspired by cpython/Lib/heapq.py,
84 ! refer to it from in depth explanations of the optimizations.
86 ! called bubble-up in the literature... but we keep cpython's name.
87 :: sift-down ( heap from to -- )
89 to data data-nth :> tmp
91 to t [ over from > and ] [
94 dup tmp heap heap-compare [
95 rot data data-set-nth t
101 tmp swap data data-set-nth ; inline
106 [ <entry> dup ] [ data>> data-push ] [ 0 rot sift-down ] tri ;
108 : heap-push ( value key heap -- )
111 : heap-push-all ( assoc heap -- )
112 '[ swap _ heap-push ] assoc-each ;
116 ! called bubble-down in the literature... but we keep cpython's name.
117 ! A quote from cpython's implementation:
118 ! > We *could* break out of the loop as soon as we find a pos where newitem <=
119 ! > both its children, but turns out that's not a good idea [...]
120 ! Indeed the code is 33% slower if we remove this optimization.
121 :: sift-up ( heap n -- )
124 n data data-nth :> tmp
126 n dup left [ dup end < ] [
129 2dup [ data data-nth ] bi@ heap heap-compare
132 [ data data-nth swap data data-set-nth ]
136 tmp over data data-set-nth
137 heap n rot sift-down ; inline
142 dup data>> f over first index<< [ pop ] keep
143 [ 2drop ] [ set-first 0 sift-up ] if-empty ;
145 : heap-pop ( heap -- value key )
146 [ heap-peek ] [ heap-pop* ] bi ;
148 : slurp-heap ( ... heap quot: ( ... value key -- ... ) -- ... )
149 [ drop '[ _ heap-empty? ] ]
150 [ '[ _ heap-pop @ ] until ] 2bi ; inline
152 : heap-pop-all ( heap -- alist )
153 [ heap-size <vector> ] keep
154 [ swap 2array suffix! ] slurp-heap { } like ;
156 ERROR: bad-heap-delete ;
158 M: bad-heap-delete summary
159 drop "Invalid entry passed to heap-delete" ;
163 : entry>index ( entry heap -- n )
164 over heap>> eq? [ bad-heap-delete ] unless
165 index>> dup [ bad-heap-delete ] unless
166 { fixnum } declare ; inline
170 M:: heap heap-delete ( entry heap -- )
171 entry heap entry>index :> n
173 data pop :> nth-entry
176 nth-entry n data data-set-nth
177 n 0 = [ t ] [ nth-entry n up data data-nth heap heap-compare ] if
178 [ heap n sift-up ] [ heap 0 n sift-down ] if
181 : >min-heap ( assoc -- min-heap )
182 dup assoc-size <vector> min-heap boa
183 [ heap-push-all ] keep ;
185 : >max-heap ( assoc -- max-heap )
186 dup assoc-size <vector> max-heap boa
187 [ heap-push-all ] keep ;