! Copyright (C) 2007, 2008 Ryan Murphy, Doug Coleman,
! Slava Pestov.
! See http://factorcode.org/license.txt for BSD license.
-USING: accessors arrays assocs combinators fry growable kernel
-kernel.private math math.order math.private sequences
-sequences.private summary vectors ;
-
+USING: accessors arrays assocs fry kernel kernel.private locals
+math math.order math.private sequences sequences.private summary
+vectors ;
IN: heaps
GENERIC: heap-push* ( value key heap -- entry )
TUPLE: heap { data vector } ;
: <heap> ( class -- heap )
- [ V{ } clone ] dip boa ; inline
+ 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 ;
: data-nth ( n heap -- entry )
data>> nth-unsafe { entry } declare ; inline
-: left-value ( n heap -- entry )
- [ left ] dip data-nth ; inline
-
-: right-value ( n heap -- entry )
- [ right ] dip data-nth ; inline
+: data-first ( heap -- entry )
+ 0 swap data-nth ; inline
: data-set-nth ( entry n heap -- )
- [ [ swap index<< ] 2keep ] dip
- data>> set-nth-unsafe ; inline
+ [ [ >>index ] keep ] dip data>> set-nth-unsafe ; inline
: data-push ( entry heap -- n )
- dup heap-size [
- swap
- [ data>> ensure 2drop ]
- [ data-set-nth ] 2bi
- ] keep ; inline
-
-: data-first ( heap -- entry )
- data>> first ; inline
-
-: data-exchange ( m n heap -- )
- [ '[ _ data-nth ] bi@ ]
- [ '[ _ data-set-nth ] bi@ ] 3bi ; inline
+ [ heap-size [ >>index ] keep ]
+ [ data>> [ set-nth ] 2keep drop ] bi ; inline
GENERIC: heap-compare ( entry1 entry2 heap -- ? )
M: max-heap heap-compare
drop { entry entry } declare [ key>> ] bi@ before? ; inline
-: heap-bounds-check? ( m heap -- ? )
- heap-size >= ; inline
+: data-compare ( m n heap -- ? )
+ [ '[ _ data-nth ] bi@ ] [ heap-compare ] bi ; inline
-: left-bounds-check? ( m heap -- ? )
- [ left ] dip heap-bounds-check? ; inline
+PRIVATE>
-: right-bounds-check? ( m heap -- ? )
- [ right ] dip heap-bounds-check? ; inline
+: >entry< ( entry -- value key )
+ [ value>> ] [ key>> ] bi ; inline
-: continue? ( m n heap -- ? )
- [ data-nth nip ]
- [ nip data-nth ]
- [ 2nip ] 3tri heap-compare ; inline
+M: heap heap-peek ( heap -- value key )
+ data-first >entry< ;
-DEFER: up-heap
+<PRIVATE
-: (up-heap) ( n heap -- )
- [ dup up ] dip
- 3dup continue? [
- [ data-exchange ] [ up-heap ] 2bi
- ] [
- 3drop
- ] if ; inline recursive
+:: sift-down ( heap from to -- )
+ to heap data-nth :> tmp
-: up-heap ( n heap -- )
- over 0 > [ (up-heap) ] [ 2drop ] if ; inline recursive
+ to t [ over from > and ] [
+ dup up
+ dup heap data-nth
+ dup tmp heap heap-compare [
+ rot heap data-set-nth t
+ ] [
+ 2drop f
+ ] if
+ ] while
-: (child) ( m heap -- n )
- { [ drop ] [ left-value ] [ right-value ] [ nip ] } 2cleave
- heap-compare [ right ] [ left ] if ; inline
+ tmp swap heap data-set-nth ; inline
-: child ( m heap -- n )
- 2dup right-bounds-check?
- [ drop left ] [ (child) ] if ; inline
+PRIVATE>
-DEFER: down-heap
+M: heap heap-push*
+ [ <entry> dup ] [ data-push ] [ 0 rot sift-down ] tri ;
-: (down-heap) ( m heap -- )
- [ drop ] [ child ] [ nip ] 2tri
- 3dup continue? [
- 3drop
- ] [
- [ data-exchange ] [ down-heap ] 2bi
- ] if ; inline recursive
+: heap-push ( value key heap -- )
+ heap-push* drop ;
-: down-heap ( m heap -- )
- 2dup left-bounds-check?
- [ 2drop ] [ (down-heap) ] if ; inline recursive
+: heap-push-all ( assoc heap -- )
+ '[ swap _ heap-push ] assoc-each ;
-PRIVATE>
+<PRIVATE
-M: heap heap-push* ( value key heap -- entry )
- [ <entry> dup ] [ data-push ] [ up-heap ] tri ;
+:: sift-up ( heap n -- )
+ heap heap-size :> end
+ n heap data-nth :> tmp
-: heap-push ( value key heap -- ) heap-push* drop ;
+ n dup left [ dup end < ] [
+ dup 1 fixnum+fast
+ dup end < [ 2dup heap data-compare ] [ f ] if
+ [ nip ] [ drop ] if
+ [ heap data-nth swap heap data-set-nth ]
+ [ dup left ] bi
+ ] while drop
-: heap-push-all ( assoc heap -- )
- '[ swap _ heap-push ] assoc-each ;
+ tmp over heap data-set-nth
+ heap n rot sift-down ; inline
-: >entry< ( entry -- value key )
- [ value>> ] [ key>> ] bi ; inline
+PRIVATE>
-M: heap heap-peek ( heap -- value key )
- data-first >entry< ;
+M: heap heap-pop*
+ dup data>> dup length 1 > [
+ [ pop ] [ set-first ] bi 0 sift-up
+ ] [
+ pop* drop
+ ] if ; inline
+
+M: heap heap-pop
+ [ data-first >entry< ] [ heap-pop* ] bi ;
+
+: heap-pop-all ( heap -- alist )
+ check-heap [ heap-size ] keep
+ '[ _ heap-pop swap 2array ] replicate ;
+
+: slurp-heap ( heap quot: ( value key -- ) -- )
+ [ check-heap [ heap-size ] keep ] dip
+ '[ _ heap-pop @ ] times ; inline
ERROR: bad-heap-delete ;
M: bad-heap-delete summary
drop "Invalid entry passed to heap-delete" ;
+<PRIVATE
+
: entry>index ( entry heap -- n )
over heap>> eq? [ bad-heap-delete ] unless
index>> { fixnum } declare ; inline
-M: heap heap-delete ( entry heap -- )
+PRIVATE>
+
+M: heap heap-delete
[ entry>index ] keep
2dup heap-size 1 - = [
nip data>> pop*
] [
[ nip data>> pop ]
[ data-set-nth ]
- [ down-heap ] 2tri
+ [ swap sift-up ] 2tri
] if ;
-M: heap heap-pop* ( heap -- )
- [ data-first ] keep heap-delete ;
-
-M: heap heap-pop ( heap -- value key )
- [ data-first dup ] keep heap-delete >entry< ;
-
-: heap-pop-all ( heap -- alist )
- [ dup heap-empty? not ]
- [ dup heap-pop swap 2array ]
- produce nip ;
-
-ERROR: not-a-heap object ;
-
-: check-heap ( heap -- heap )
- dup heap? [ not-a-heap ] unless ; inline
-
-: slurp-heap ( heap quot: ( value key -- ) -- )
- [ check-heap ] dip over heap-empty? [ 2drop ] [
- [ [ heap-pop ] dip call ] [ slurp-heap ] 2bi
- ] if ; inline recursive
-
: >min-heap ( assoc -- min-heap )
<min-heap> [ heap-push-all ] keep ;