! Copyright (C) 2007, 2008 Ryan Murphy, Doug Coleman,
! Slava Pestov.
! See http://factorcode.org/license.txt for BSD license.
-USING: accessors arrays assocs fry kernel kernel.private locals
+USING: accessors arrays assocs fry kernel kernel.private
math math.order math.private sequences sequences.private summary
vectors ;
IN: heaps
GENERIC: heap-push* ( value key heap -- entry )
GENERIC: heap-peek ( heap -- value key )
GENERIC: heap-pop* ( heap -- )
-GENERIC: heap-pop ( heap -- value key )
GENERIC: heap-delete ( entry heap -- )
GENERIC: heap-empty? ( heap -- ? )
GENERIC: heap-size ( heap -- n )
: <heap> ( class -- heap )
V{ } clone swap boa ; inline
-ERROR: not-a-heap object ;
-
-: check-heap ( heap -- heap )
- dup heap? [ not-a-heap ] unless ; inline
-
TUPLE: entry value key heap index ;
: <entry> ( value key heap -- entry )
: <max-heap> ( -- max-heap ) max-heap <heap> ;
-M: heap heap-empty? ( heap -- ? )
- data>> empty? ; inline
+M: heap heap-empty? data>> empty? ; inline
-M: heap heap-size ( heap -- n )
- data>> length ; inline
+M: heap heap-size data>> length ; inline
<PRIVATE
: up ( n -- m )
{ fixnum } declare 1 fixnum-fast 2/ ; inline
-: data-nth ( n heap -- entry )
- data>> nth-unsafe { entry } declare ; inline
-
-: data-first ( heap -- entry )
- data>> first ; inline
+: data-nth ( n data -- entry )
+ nth-unsafe { entry } declare ; inline
-: data-set-nth ( entry n heap -- )
- [ [ >>index ] keep ] dip data>> set-nth-unsafe ; inline
+: data-set-nth ( entry n data -- )
+ [ [ >>index ] keep ] dip set-nth-unsafe ; inline
-: data-push ( entry heap -- n )
- [ heap-size [ >>index ] keep ]
- [ data>> [ set-nth ] 2keep drop ] bi ; inline
+: data-push ( entry data -- n )
+ [ length [ >>index ] keep ]
+ [ [ set-nth ] keepd ] bi ; inline
GENERIC: heap-compare ( entry1 entry2 heap -- ? )
M: max-heap heap-compare
drop { entry entry } declare [ key>> ] bi@ before? ; inline
-: data-compare ( m n heap -- ? )
- [ '[ _ data-nth ] bi@ ] [ heap-compare ] bi ; inline
-
PRIVATE>
: >entry< ( entry -- value key )
[ value>> ] [ key>> ] bi ; inline
-M: heap heap-peek ( heap -- value key )
- data-first >entry< ;
+M: heap heap-peek
+ data>> first >entry< ;
<PRIVATE
+! names and optimizations inspired by cpython/Lib/heapq.py,
+! refer to it from in depth explanations of the optimizations.
+
+! called bubble-up in the literature... but we keep cpython's name.
:: sift-down ( heap from to -- )
- to heap data-nth :> tmp
+ heap data>> :> data
+ to data data-nth :> tmp
to t [ over from > and ] [
dup up
- dup heap data-nth
+ dup data data-nth
dup tmp heap heap-compare [
- rot heap data-set-nth t
+ rot data data-set-nth t
] [
2drop f
] if
] while
- tmp swap heap data-set-nth ; inline
+ tmp swap data data-set-nth ; inline
PRIVATE>
M: heap heap-push*
- [ <entry> dup ] [ data-push ] [ 0 rot sift-down ] tri ;
+ [ <entry> dup ] [ data>> data-push ] [ 0 rot sift-down ] tri ;
: heap-push ( value key heap -- )
heap-push* drop ;
<PRIVATE
+! called bubble-down in the literature... but we keep cpython's name.
+! A quote from cpython's implementation:
+! > We *could* break out of the loop as soon as we find a pos where newitem <=
+! > both its children, but turns out that's not a good idea [...]
+! Indeed the code is 33% slower if we remove this optimization.
:: sift-up ( heap n -- )
- heap heap-size :> end
- n heap data-nth :> tmp
+ heap data>> :> data
+ data length :> end
+ n data data-nth :> tmp
n dup left [ dup end < ] [
dup 1 fixnum+fast
- dup end < [ 2dup heap data-compare ] [ f ] if
+ dup end < [
+ 2dup [ data data-nth ] bi@ heap heap-compare
+ ] [ f ] if
[ nip ] [ drop ] if
- [ heap data-nth swap heap data-set-nth ]
+ [ data data-nth swap data data-set-nth ]
[ dup left ] bi
] while drop
- tmp over heap data-set-nth
+ tmp over data data-set-nth
heap n rot sift-down ; inline
PRIVATE>
M: heap heap-pop*
- dup data>> dup length 1 > [
- [ pop ] [ set-first ] bi 0 sift-up
- ] [
- pop* drop
- ] if ; inline
+ dup data>> f over first index<< [ pop ] keep
+ [ 2drop ] [ set-first 0 sift-up ] if-empty ;
-M: heap heap-pop
- [ data-first >entry< ] [ heap-pop* ] bi ;
+: heap-pop ( heap -- value key )
+ [ heap-peek ] [ heap-pop* ] bi ;
: slurp-heap ( ... heap quot: ( ... value key -- ... ) -- ... )
- [ check-heap ] dip
[ drop '[ _ heap-empty? ] ]
[ '[ _ heap-pop @ ] until ] 2bi ; inline
: entry>index ( entry heap -- n )
over heap>> eq? [ bad-heap-delete ] unless
- index>> { fixnum } declare ; inline
+ index>> dup [ bad-heap-delete ] unless
+ { fixnum } declare ; inline
PRIVATE>
-M: heap heap-delete
- [ entry>index ] keep
- 2dup heap-size 1 - = [
- nip data>> pop*
- ] [
- [ nip data>> pop ]
- [ data-set-nth ]
- [ swap sift-up ] 2tri
- ] if ;
+M:: heap heap-delete ( entry heap -- )
+ entry heap entry>index :> n
+ heap data>> :> data
+ data pop :> nth-entry
+ f entry index<<
+ n data length = [
+ nth-entry n data data-set-nth
+ n 0 = [ t ] [ nth-entry n up data data-nth heap heap-compare ] if
+ [ heap n sift-up ] [ heap 0 n sift-down ] if
+ ] unless ;
: >min-heap ( assoc -- min-heap )
- <min-heap> [ heap-push-all ] keep ;
+ dup assoc-size <vector> min-heap boa
+ [ heap-push-all ] keep ;
: >max-heap ( assoc -- max-heap )
- <max-heap> [ heap-push-all ] keep ;
+ dup assoc-size <vector> max-heap boa
+ [ heap-push-all ] keep ;