Minor cleanups in trees

Doug Coleman 2010-10-09 11:21:11 -07:00
parent 97ea36e24b
commit 79b082b728
2 changed files with 20 additions and 20 deletions

View File

@ -3,13 +3,13 @@
USING: combinators kernel generic math math.functions USING: combinators kernel generic math math.functions
math.parser namespaces io sequences trees shuffle math.parser namespaces io sequences trees shuffle
assocs parser accessors math.order prettyprint.custom assocs parser accessors math.order prettyprint.custom
trees.private ; trees.private fry ;
IN: trees.avl IN: trees.avl
TUPLE: avl < tree ; TUPLE: avl < tree ;
: <avl> ( -- tree ) : <avl> ( -- tree )
avl new-tree ; avl new-tree ; inline
<PRIVATE <PRIVATE
@ -17,15 +17,16 @@ TUPLE: avl-node < node balance ;
: <avl-node> ( key value -- node ) : <avl-node> ( key value -- node )
avl-node new-node avl-node new-node
0 >>balance ; 0 >>balance ; inline
: increase-balance ( node amount -- ) : increase-balance ( node amount -- node )
swap [ + ] change-balance drop ; '[ _ + ] change-balance ;
: rotate ( node -- node ) : rotate ( node -- node )
dup node+link dup
dup node-link [ node+link ]
pick set-node+link [ node-link ]
[ set-node+link ] tri
[ set-node-link ] keep ; [ set-node-link ] keep ;
: single-rotate ( node -- node ) : single-rotate ( node -- node )
@ -36,8 +37,8 @@ TUPLE: avl-node < node balance ;
: pick-balances ( a node -- balance balance ) : pick-balances ( a node -- balance balance )
balance>> { balance>> {
{ [ dup zero? ] [ 2drop 0 0 ] } { [ dup zero? ] [ 2drop 0 0 ] }
{ [ over = ] [ neg 0 ] } { [ 2dup = ] [ nip neg 0 ] }
[ 0 swap ] [ drop 0 swap ]
} cond ; } cond ;
: double-rotate ( node -- node ) : double-rotate ( node -- node )
@ -57,9 +58,8 @@ TUPLE: avl-node < node balance ;
: balance-insert ( node -- node taller? ) : balance-insert ( node -- node taller? )
dup balance>> { dup balance>> {
{ [ dup zero? ] [ drop f ] } { [ dup zero? ] [ drop f ] }
{ [ dup abs 2 = ] { [ dup abs 2 = ] [ sgn neg [ select-rotate ] with-side f ] }
[ sgn neg [ select-rotate ] with-side f ] } [ drop t ] ! balance is -1 or 1, tree is taller
{ [ drop t ] [ t ] } ! balance is -1 or 1, tree is taller
} cond ; } cond ;
DEFER: avl-set DEFER: avl-set
@ -68,7 +68,7 @@ DEFER: avl-set
2dup key>> before? left right ? [ 2dup key>> before? left right ? [
[ node-link avl-set ] keep swap [ node-link avl-set ] keep swap
[ [ set-node-link ] keep ] dip [ [ set-node-link ] keep ] dip
[ dup current-side get increase-balance balance-insert ] [ current-side get increase-balance balance-insert ]
[ f ] if [ f ] if
] with-side ; ] with-side ;
@ -95,14 +95,14 @@ M: avl set-at ( value key node -- node )
dup balance>> { dup balance>> {
{ [ dup zero? ] [ drop t ] } { [ dup zero? ] [ drop t ] }
{ [ dup abs 2 = ] [ sgn neg [ delete-select-rotate ] with-side ] } { [ dup abs 2 = ] [ sgn neg [ delete-select-rotate ] with-side ] }
{ [ drop t ] [ f ] } ! balance is -1 or 1, tree is not shorter [ drop f ] ! balance is -1 or 1, tree is not shorter
} cond ; } cond ;
: balance-delete ( node -- node shorter? ) : balance-delete ( node -- node shorter? )
current-side get over balance>> { current-side get over balance>> {
{ [ dup zero? ] [ drop neg over balance<< f ] } { [ dup zero? ] [ drop neg over balance<< f ] }
{ [ dupd = ] [ drop 0 >>balance t ] } { [ 2dup = ] [ 2drop 0 >>balance t ] }
[ dupd neg increase-balance rebalance-delete ] [ drop neg increase-balance rebalance-delete ]
} cond ; } cond ;
: avl-replace-with-extremity ( to-replace node -- node shorter? ) : avl-replace-with-extremity ( to-replace node -- node shorter? )
@ -155,7 +155,7 @@ M: avl new-assoc 2drop <avl> ;
PRIVATE> PRIVATE>
: >avl ( assoc -- avl ) : >avl ( assoc -- avl )
T{ avl f f 0 } assoc-clone-like ; T{ avl } assoc-clone-like ;
M: avl assoc-like M: avl assoc-like
drop dup avl? [ >avl ] unless ; drop dup avl? [ >avl ] unless ;

View File

@ -5,7 +5,7 @@ prettyprint.private kernel.private assocs random combinators
parser math.order accessors deques make prettyprint.custom ; parser math.order accessors deques make prettyprint.custom ;
IN: trees IN: trees
TUPLE: tree root count ; TUPLE: tree root { count integer } ;
<PRIVATE <PRIVATE
@ -28,7 +28,7 @@ TUPLE: node key value left right ;
: new-node ( key value class -- node ) : new-node ( key value class -- node )
new new
swap >>value swap >>value
swap >>key ; swap >>key ; inline
: <node> ( key value -- node ) : <node> ( key value -- node )
node new-node ; node new-node ;