! See http://factorcode.org/license.txt for BSD license.
USING: arrays kernel math namespaces tools.test
heaps heaps.private math.parser random assocs sequences sorting
-accessors math.order ;
+accessors math.order locals ;
IN: heaps.tests
[ <min-heap> heap-pop ] must-fail
[ 0 ] [ <max-heap> heap-size ] unit-test
[ 1 ] [ <max-heap> t 1 pick heap-push heap-size ] unit-test
-: heap-sort ( alist -- keys )
- <min-heap> [ heap-push-all ] keep heap-pop-all ;
+: heap-sort ( alist heap -- keys )
+ [ heap-push-all ] keep heap-pop-all ;
: random-alist ( n -- alist )
iota [
drop 32 random-bits dup number>string
- ] H{ } map>assoc ;
+ ] H{ } map>assoc >alist ;
-: test-heap-sort ( n -- ? )
- random-alist dup >alist sort-keys swap heap-sort = ;
+:: test-heap-sort ( n heap reverse? -- ? )
+ n random-alist
+ [ sort-keys reverse? [ reverse ] when ] keep
+ heap heap-sort = ;
+
+: test-minheap-sort ( n -- ? )
+ <min-heap> f test-heap-sort ;
+
+: test-maxheap-sort ( n -- ? )
+ <max-heap> t test-heap-sort ;
+
+14 [
+ [ t ] swap [ 2^ <min-heap> f test-heap-sort ] curry unit-test
+] each-integer
14 [
- [ t ] swap [ 2^ test-heap-sort ] curry unit-test
+ [ t ] swap [ 2^ <max-heap> t test-heap-sort ] curry unit-test
] each-integer
: test-entry-indices ( n -- ? )
! Slava Pestov.
! See http://factorcode.org/license.txt for BSD license.
USING: kernel math sequences arrays assocs sequences.private
-growable accessors math.order summary vectors ;
+growable accessors math.order summary vectors fry combinators ;
IN: heaps
GENERIC: heap-push* ( value key heap -- entry )
[ right ] dip data-nth ; inline
: data-set-nth ( entry n heap -- )
- [ [ >>index drop ] 2keep ] dip
+ [ [ >>index drop ] [ ] 2bi ] dip
data>> set-nth-unsafe ; inline
: data-push ( entry heap -- n )
dup heap-size [
swap 2dup data>> ensure 2drop data-set-nth
- ] keep ; inline
-
-: data-pop ( heap -- entry )
- data>> pop ; inline
-
-: data-pop* ( heap -- )
- data>> pop* ; inline
+ ] [
+ ] bi ; inline
: data-first ( heap -- entry )
data>> first ; inline
: data-exchange ( m n heap -- )
- [ [ data-nth ] curry bi@ ]
- [ [ data-set-nth ] curry bi@ ] 3bi ; inline
+ [ '[ _ data-nth ] bi@ ]
+ [ '[ _ data-set-nth ] bi@ ] 3bi ; inline
-GENERIC: heap-compare ( pair1 pair2 heap -- ? )
+GENERIC: heap-compare ( entry1 entry2 heap -- ? )
-: (heap-compare) ( pair1 pair2 heap -- <=> )
+: (heap-compare) ( entry1 entry2 heap -- <=> )
drop [ key>> ] compare ; inline
M: min-heap heap-compare (heap-compare) +gt+ eq? ;
: right-bounds-check? ( m heap -- ? )
[ right ] dip heap-bounds-check? ; inline
-: continue? ( m up[m] heap -- ? )
- [ data-nth swap ] keep [ data-nth ] keep
- heap-compare ; inline
+: continue? ( m n heap -- ? )
+ [ data-nth nip ]
+ [ nip data-nth ]
+ [ 2nip ] 3tri heap-compare ;
DEFER: up-heap
: (up-heap) ( n heap -- )
[ dup up ] dip
3dup continue? [
- [ data-exchange ] 2keep up-heap
+ [ data-exchange ] [ up-heap ] 2bi
] [
3drop
] if ; inline recursive
over 0 > [ (up-heap) ] [ 2drop ] if ; inline recursive
: (child) ( m heap -- n )
- 2dup right-value
- [ 2dup left-value ] dip
- rot heap-compare
- [ right ] [ left ] if ;
+ { [ drop ] [ left-value ] [ right-value ] [ nip ] } 2cleave
+ heap-compare [ right ] [ left ] if ;
: child ( m heap -- n )
2dup right-bounds-check?
DEFER: down-heap
: (down-heap) ( m heap -- )
- [ child ] 2keep swapd
+ [ drop ] [ child ] [ nip ] 2tri
3dup continue? [
3drop
] [
- [ data-exchange ] 2keep down-heap
+ [ data-exchange ] [ down-heap ] 2bi
] if ; inline recursive
: down-heap ( m heap -- )
PRIVATE>
M: heap heap-push* ( value key heap -- entry )
- [ <entry> dup ] keep [ data-push ] keep up-heap ;
+ [ <entry> dup ] [ data-push ] [ ] tri up-heap ;
: heap-push ( value key heap -- ) heap-push* drop ;
: heap-push-all ( assoc heap -- )
- [ swapd heap-push ] curry assoc-each ;
+ '[ swap _ heap-push ] assoc-each ;
-: >entry< ( entry -- key value )
+: >entry< ( entry -- value key )
[ value>> ] [ key>> ] bi ; inline
M: heap heap-peek ( heap -- value key )
index>> ;
M: heap heap-delete ( entry heap -- )
- [ entry>index ] keep
+ [ entry>index ] [ ] bi
2dup heap-size 1 - = [
- nip data-pop*
+ nip data>> pop*
] [
- [ nip data-pop ] 2keep
- [ data-set-nth ] 2keep
+ [ nip data>> pop ]
+ [ data-set-nth ]
+ [ ] 2tri
down-heap
] if ;
M: heap heap-pop* ( heap -- )
- dup data-first swap heap-delete ;
+ [ data-first ] keep heap-delete ;
M: heap heap-pop ( heap -- value key )
- dup data-first [ swap heap-delete ] keep >entry< ;
+ [ data-first ] keep
+ [ heap-delete ] [ drop ] 2bi >entry< ;
: heap-pop-all ( heap -- alist )
[ dup heap-empty? not ]
[ 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