trees.avl, implement assoc-size

char-rename
Jon Harper 2017-01-06 17:33:05 +01:00 committed by John Benediktsson
parent d2cfbafa13
commit 4442d922b0
2 changed files with 17 additions and 11 deletions

View File

@ -119,3 +119,7 @@ IN: trees.avl.tests
{ f } [ test-tree 9 over delete-at 9 of ] unit-test { f } [ test-tree 9 over delete-at 9 of ] unit-test
{ "replaced seven" } [ test-tree 9 over delete-at 7 of ] unit-test { "replaced seven" } [ test-tree 9 over delete-at 7 of ] unit-test
{ "nine" } [ test-tree 7 over delete-at 4 over delete-at 9 of ] unit-test { "nine" } [ test-tree 7 over delete-at 4 over delete-at 9 of ] unit-test
! test assoc-size
{ 3 } [ test-tree assoc-size ] unit-test
{ 2 } [ test-tree 9 over delete-at assoc-size ] unit-test

View File

@ -63,24 +63,25 @@ TUPLE: avl-node < node balance ;
DEFER: avl-set DEFER: avl-set
: avl-insert ( value key node -- node taller? ) : avl-insert ( value key node -- node taller? created? )
2dup key>> before? left right ? [ 2dup key>> before? left right ? [
[ node-link avl-set ] keep swap [ node-link avl-set ] keep -rot
[ [ set-node-link ] keep ] dip [ [ set-node-link ] keep ] 2dip swap
[ current-side get increase-balance balance-insert ] [ [ current-side get increase-balance balance-insert ] dip ]
[ f ] if [ f swap ] if
] with-side ; ] with-side ;
: (avl-set) ( value key node -- node taller? ) : (avl-set) ( value key node -- node taller? created? )
2dup key>> = [ 2dup key>> = [
-rot pick key<< >>value f -rot pick key<< >>value f f
] [ avl-insert ] if ; ] [ avl-insert ] if ;
: avl-set ( value key node -- node taller? ) : avl-set ( value key node -- node taller? created? )
[ (avl-set) ] [ swap <avl-node> t ] if* ; [ (avl-set) ] [ swap <avl-node> t t ] if* ;
M: avl set-at ( value key node -- ) M: avl set-at ( value key node -- )
[ avl-set drop ] change-root drop ; [ avl-set nip swap ] change-root
swap [ dup inc-count ] when drop ;
: delete-select-rotate ( node -- node shorter? ) : delete-select-rotate ( node -- node shorter? )
dup node+link balance>> zero? [ dup node+link balance>> zero? [
@ -147,7 +148,8 @@ M: avl-node avl-delete ( key node -- node shorter? deleted? )
] if-zero ; ] if-zero ;
M: avl delete-at ( key node -- ) M: avl delete-at ( key node -- )
[ avl-delete 2drop ] change-root drop ; [ avl-delete nip swap ] change-root
swap [ dup dec-count ] when drop ;
M: avl new-assoc 2drop <avl> ; M: avl new-assoc 2drop <avl> ;