! Slava Pestov.
! See http://factorcode.org/license.txt for BSD license.
USING: kernel math sequences arrays assocs sequences.private
-growable accessors math.order summary ;
+growable accessors math.order summary vectors ;
IN: heaps
GENERIC: heap-push* ( value key heap -- entry )
<PRIVATE
-TUPLE: heap data ;
+TUPLE: heap { data vector } ;
: <heap> ( class -- heap )
[ V{ } clone ] dip boa ; inline
TUPLE: entry value key heap index ;
-: <entry> ( value key heap -- entry ) f entry boa ;
+: <entry> ( value key heap -- entry ) f entry boa ; inline
PRIVATE>
: right ( n -- m ) 1 shift 2 + ; inline
-: up ( n -- m ) 1- 2/ ; inline
+: up ( n -- m ) 1 - 2/ ; inline
: data-nth ( n heap -- entry )
data>> nth-unsafe ; inline
-: up-value ( n heap -- entry )
- [ up ] dip data-nth ; inline
-
: left-value ( n heap -- entry )
[ left ] dip data-nth ; inline
: data-pop* ( heap -- )
data>> pop* ; inline
-: data-peek ( heap -- entry )
- data>> peek ; inline
-
: data-first ( heap -- entry )
data>> first ; inline
[ data-exchange ] 2keep up-heap
] [
3drop
- ] if ;
+ ] if ; inline recursive
: up-heap ( n heap -- )
- over 0 > [ (up-heap) ] [ 2drop ] if ;
+ over 0 > [ (up-heap) ] [ 2drop ] if ; inline recursive
: (child) ( m heap -- n )
2dup right-value
2dup right-bounds-check?
[ drop left ] [ (child) ] if ;
-: swap-down ( m heap -- )
- [ child ] 2keep data-exchange ;
-
DEFER: down-heap
: (down-heap) ( m heap -- )
3drop
] [
[ data-exchange ] 2keep down-heap
- ] if ;
+ ] if ; inline recursive
: down-heap ( m heap -- )
- 2dup left-bounds-check? [ 2drop ] [ (down-heap) ] if ;
+ 2dup left-bounds-check? [ 2drop ] [ (down-heap) ] if ; inline recursive
PRIVATE>
[ swapd heap-push ] curry assoc-each ;
: >entry< ( entry -- key value )
- [ value>> ] [ key>> ] bi ;
+ [ value>> ] [ key>> ] bi ; inline
M: heap heap-peek ( heap -- value key )
data-first >entry< ;
M: heap heap-delete ( entry heap -- )
[ entry>index ] keep
- 2dup heap-size 1- = [
+ 2dup heap-size 1 - = [
nip data-pop*
] [
[ nip data-pop ] 2keep
[ dup heap-pop swap 2array ]
produce nip ;
+: heap-values ( heap -- alist )
+ data>> [ value>> ] { } map-as ;
+
: slurp-heap ( heap quot: ( elt -- ) -- )
over heap-empty? [ 2drop ] [
[ [ heap-pop drop ] dip call ] [ slurp-heap ] 2bi