--- /dev/null
+Alex Chapman
+Daniel Ehrenberg
--- /dev/null
+Alex Chapman
+Daniel Ehrenberg
--- /dev/null
+USING: help.syntax help.markup assocs ;
+IN: trees.avl
+
+HELP: AVL{
+{ $syntax "AVL{ { key value }... }" }
+{ $values { "key" "a key" } { "value" "a value" } }
+{ $description "Literal syntax for an AVL tree." } ;
+
+HELP: <avl>
+{ $values { "tree" avl } }
+{ $description "Creates an empty AVL tree" } ;
+
+HELP: >avl
+{ $values { "assoc" assoc } { "avl" avl } }
+{ $description "Converts any " { $link assoc } " into an AVL tree." } ;
+
+HELP: avl
+{ $class-description "This is the class for AVL trees. These conform to the assoc protocol and have efficient (logarithmic time) storage and retrieval operations." } ;
+
+ARTICLE: "trees.avl" "AVL trees"
+"This is a library for AVL trees, with logarithmic time storage and retrieval operations. These trees conform to the assoc protocol."
+{ $subsection avl }
+{ $subsection <avl> }
+{ $subsection >avl }
+{ $subsection POSTPONE: AVL{ } ;
+
+ABOUT: "trees.avl"
--- /dev/null
+USING: kernel tools.test trees trees.avl math random sequences
+assocs accessors ;
+IN: trees.avl.tests
+
+[ "key1" 0 "key2" 0 ] [
+ T{ avl-node f "key1" f f T{ avl-node f "key2" f f 1 } 2 }
+ [ single-rotate ] go-left
+ [ left>> dup key>> swap balance>> ] keep
+ dup key>> swap balance>>
+] unit-test
+
+[ "key1" 0 "key2" 0 ] [
+ T{ avl-node f "key1" f f T{ avl-node f "key2" f f f 1 } 2 }
+ [ select-rotate ] go-left
+ [ left>> dup key>> swap balance>> ] keep
+ dup key>> swap balance>>
+] unit-test
+
+[ "key1" 0 "key2" 0 ] [
+ T{ avl-node f "key1" f T{ avl-node f "key2" f f f -1 } f -2 }
+ [ single-rotate ] go-right
+ [ right>> dup key>> swap balance>> ] keep
+ dup key>> swap balance>>
+] unit-test
+
+[ "key1" 0 "key2" 0 ] [
+ T{ avl-node f "key1" f T{ avl-node f "key2" f f f -1 } f -2 }
+ [ select-rotate ] go-right
+ [ right>> dup key>> swap balance>> ] keep
+ dup key>> swap balance>>
+] unit-test
+
+[ "key1" -1 "key2" 0 "key3" 0 ]
+[ T{ avl-node f "key1" f f
+ T{ avl-node f "key2" f
+ T{ avl-node f "key3" f f f 1 } f -1 } 2 }
+ [ double-rotate ] go-left
+ [ left>> dup key>> swap balance>> ] keep
+ [ right>> dup key>> swap balance>> ] keep
+ dup key>> swap balance>> ] unit-test
+[ "key1" 0 "key2" 0 "key3" 0 ]
+[ T{ avl-node f "key1" f f
+ T{ avl-node f "key2" f
+ T{ avl-node f "key3" f f f 0 } f -1 } 2 }
+ [ double-rotate ] go-left
+ [ left>> dup key>> swap balance>> ] keep
+ [ right>> dup key>> swap balance>> ] keep
+ dup key>> swap balance>> ] unit-test
+[ "key1" 0 "key2" 1 "key3" 0 ]
+[ T{ avl-node f "key1" f f
+ T{ avl-node f "key2" f
+ T{ avl-node f "key3" f f f -1 } f -1 } 2 }
+ [ double-rotate ] go-left
+ [ left>> dup key>> swap balance>> ] keep
+ [ right>> dup key>> swap balance>> ] keep
+ dup key>> swap balance>> ] unit-test
+
+[ "key1" 1 "key2" 0 "key3" 0 ]
+[ T{ avl-node f "key1" f
+ T{ avl-node f "key2" f f
+ T{ avl-node f "key3" f f f -1 } 1 } f -2 }
+ [ double-rotate ] go-right
+ [ right>> dup key>> swap balance>> ] keep
+ [ left>> dup key>> swap balance>> ] keep
+ dup key>> swap balance>> ] unit-test
+[ "key1" 0 "key2" 0 "key3" 0 ]
+[ T{ avl-node f "key1" f
+ T{ avl-node f "key2" f f
+ T{ avl-node f "key3" f f f 0 } 1 } f -2 }
+ [ double-rotate ] go-right
+ [ right>> dup key>> swap balance>> ] keep
+ [ left>> dup key>> swap balance>> ] keep
+ dup key>> swap balance>> ] unit-test
+[ "key1" 0 "key2" -1 "key3" 0 ]
+[ T{ avl-node f "key1" f
+ T{ avl-node f "key2" f f
+ T{ avl-node f "key3" f f f 1 } 1 } f -2 }
+ [ double-rotate ] go-right
+ [ right>> dup key>> swap balance>> ] keep
+ [ left>> dup key>> swap balance>> ] keep
+ dup key>> swap balance>> ] unit-test
+
+[ "eight" ] [
+ <avl> "seven" 7 pick set-at
+ "eight" 8 pick set-at "nine" 9 pick set-at
+ root>> value>>
+] unit-test
+
+[ "another eight" ] [ ! ERROR!
+ <avl> "seven" 7 pick set-at
+ "another eight" 8 pick set-at 8 swap at
+] unit-test
+
+: test-tree ( -- tree )
+ AVL{
+ { 7 "seven" }
+ { 9 "nine" }
+ { 4 "four" }
+ { 4 "replaced four" }
+ { 7 "replaced seven" }
+ } clone ;
+
+! test set-at, at, at*
+[ t ] [ test-tree avl? ] unit-test
+[ "seven" ] [ <avl> "seven" 7 pick set-at 7 swap at ] unit-test
+[ "seven" t ] [ <avl> "seven" 7 pick set-at 7 swap at* ] unit-test
+[ f f ] [ <avl> "seven" 7 pick set-at 8 swap at* ] unit-test
+[ "seven" ] [ <avl> "seven" 7 pick set-at 7 swap at ] unit-test
+[ "replacement" ] [ <avl> "seven" 7 pick set-at "replacement" 7 pick set-at 7 swap at ] unit-test
+[ "nine" ] [ test-tree 9 swap at ] unit-test
+[ "replaced four" ] [ test-tree 4 swap at ] unit-test
+[ "replaced seven" ] [ test-tree 7 swap at ] unit-test
+
+! test delete-at--all errors!
+[ f ] [ test-tree 9 over delete-at 9 swap at ] unit-test
+[ "replaced seven" ] [ test-tree 9 over delete-at 7 swap at ] unit-test
+[ "nine" ] [ test-tree 7 over delete-at 4 over delete-at 9 swap at ] unit-test
--- /dev/null
+! Copyright (C) 2007 Alex Chapman
+! See http://factorcode.org/license.txt for BSD license.
+USING: combinators kernel generic math math.functions
+math.parser namespaces io prettyprint.backend sequences trees
+assocs parser accessors math.order ;
+IN: trees.avl
+
+TUPLE: avl < tree ;
+
+: <avl> ( -- tree )
+ avl new-tree ;
+
+TUPLE: avl-node < node balance ;
+
+: <avl-node> ( key value -- node )
+ avl-node new-node
+ 0 >>balance ;
+
+: increase-balance ( node amount -- )
+ swap [ + ] change-balance drop ;
+
+: rotate ( node -- node )
+ dup node+link dup node-link pick set-node+link
+ tuck set-node-link ;
+
+: single-rotate ( node -- node )
+ 0 over (>>balance) 0 over node+link
+ (>>balance) rotate ;
+
+: pick-balances ( a node -- balance balance )
+ balance>> {
+ { [ dup zero? ] [ 2drop 0 0 ] }
+ { [ over = ] [ neg 0 ] }
+ [ 0 swap ]
+ } cond ;
+
+: double-rotate ( node -- node )
+ [
+ node+link [
+ node-link current-side get neg
+ over pick-balances rot 0 swap (>>balance)
+ ] keep (>>balance)
+ ] keep swap >>balance
+ dup node+link [ rotate ] with-other-side
+ over set-node+link rotate ;
+
+: select-rotate ( node -- node )
+ dup node+link balance>> current-side get =
+ [ double-rotate ] [ single-rotate ] if ;
+
+: balance-insert ( node -- node taller? )
+ dup balance>> {
+ { [ dup zero? ] [ drop f ] }
+ { [ dup abs 2 = ]
+ [ sgn neg [ select-rotate ] with-side f ] }
+ { [ drop t ] [ t ] } ! balance is -1 or 1, tree is taller
+ } cond ;
+
+DEFER: avl-set
+
+: avl-insert ( value key node -- node taller? )
+ 2dup key>> before? left right ? [
+ [ node-link avl-set ] keep swap
+ [ tuck set-node-link ] dip
+ [ dup current-side get increase-balance balance-insert ]
+ [ f ] if
+ ] with-side ;
+
+: (avl-set) ( value key node -- node taller? )
+ 2dup key>> = [
+ -rot pick (>>key) over (>>value) f
+ ] [ avl-insert ] if ;
+
+: avl-set ( value key node -- node taller? )
+ [ (avl-set) ] [ swap <avl-node> t ] if* ;
+
+M: avl set-at ( value key node -- node )
+ [ avl-set drop ] change-root drop ;
+
+: delete-select-rotate ( node -- node shorter? )
+ dup node+link balance>> zero? [
+ current-side get neg over (>>balance)
+ current-side get over node+link (>>balance) rotate f
+ ] [
+ select-rotate t
+ ] if ;
+
+: rebalance-delete ( node -- node shorter? )
+ dup balance>> {
+ { [ dup zero? ] [ drop t ] }
+ { [ dup abs 2 = ] [ sgn neg [ delete-select-rotate ] with-side ] }
+ { [ drop t ] [ f ] } ! balance is -1 or 1, tree is not shorter
+ } cond ;
+
+: balance-delete ( node -- node shorter? )
+ current-side get over balance>> {
+ { [ dup zero? ] [ drop neg over (>>balance) f ] }
+ { [ dupd = ] [ drop 0 >>balance t ] }
+ [ dupd neg increase-balance rebalance-delete ]
+ } cond ;
+
+: avl-replace-with-extremity ( to-replace node -- node shorter? )
+ dup node-link [
+ swapd avl-replace-with-extremity [ over set-node-link ] dip
+ [ balance-delete ] [ f ] if
+ ] [
+ [ copy-node-contents drop ] keep node+link t
+ ] if* ;
+
+: replace-with-a-child ( node -- node shorter? )
+ #! assumes that node is not a leaf, otherwise will recurse forever
+ dup node-link [
+ dupd [ avl-replace-with-extremity ] with-other-side
+ [ over set-node-link ] dip [ balance-delete ] [ f ] if
+ ] [
+ [ replace-with-a-child ] with-other-side
+ ] if* ;
+
+: avl-delete-node ( node -- node shorter? )
+ #! delete this node, returning its replacement, and whether this subtree is
+ #! shorter as a result
+ dup leaf? [
+ drop f t
+ ] [
+ left [ replace-with-a-child ] with-side
+ ] if ;
+
+GENERIC: avl-delete ( key node -- node shorter? deleted? )
+
+M: f avl-delete ( key f -- f f f ) nip f f ;
+
+: (avl-delete) ( key node -- node shorter? deleted? )
+ tuck node-link avl-delete [
+ [ over set-node-link ] dip [ balance-delete ] [ f ] if
+ ] dip ;
+
+M: avl-node avl-delete ( key node -- node shorter? deleted? )
+ 2dup key>> key-side dup zero? [
+ drop nip avl-delete-node t
+ ] [
+ [ (avl-delete) ] with-side
+ ] if ;
+
+M: avl delete-at ( key node -- )
+ [ avl-delete 2drop ] change-root drop ;
+
+M: avl new-assoc 2drop <avl> ;
+
+: >avl ( assoc -- avl )
+ T{ avl f f 0 } assoc-clone-like ;
+
+M: avl assoc-like
+ drop dup avl? [ >avl ] unless ;
+
+: AVL{
+ \ } [ >avl ] parse-literal ; parsing
+
+! M: avl pprint-delims drop \ AVL{ \ } ;
--- /dev/null
+Balanced AVL trees
--- /dev/null
+collections
--- /dev/null
+Mackenzie Straight
+Daniel Ehrenberg
--- /dev/null
+USING: help.syntax help.markup assocs ;
+IN: trees.splay
+
+HELP: SPLAY{
+{ $syntax "SPLAY{ { key value }... }" }
+{ $values { "key" "a key" } { "value" "a value" } }
+{ $description "Literal syntax for an splay tree." } ;
+
+HELP: <splay>
+{ $values { "tree" splay } }
+{ $description "Creates an empty splay tree" } ;
+
+HELP: >splay
+{ $values { "assoc" assoc } { "tree" splay } }
+{ $description "Converts any " { $link assoc } " into an splay tree." } ;
+
+HELP: splay
+{ $class-description "This is the class for splay trees. Splay trees have amortized average-case logarithmic time storage and retrieval operations, and better complexity on more skewed lookup distributions, though in bad situations they can degrade to linear time, resembling a linked list. These conform to the assoc protocol." } ;
+
+ARTICLE: "trees.splay" "Splay trees"
+"This is a library for splay trees. Splay trees have amortized average-case logarithmic time storage and retrieval operations, and better complexity on more skewed lookup distributions, though in bad situations they can degrade to linear time, resembling a linked list. These trees conform to the assoc protocol."
+{ $subsection splay }
+{ $subsection <splay> }
+{ $subsection >splay }
+{ $subsection POSTPONE: SPLAY{ } ;
+
+ABOUT: "trees.splay"
--- /dev/null
+! Copyright (c) 2005 Mackenzie Straight.
+! See http://factorcode.org/license.txt for BSD license.
+USING: kernel tools.test trees.splay math namespaces assocs
+sequences random sets make grouping ;
+IN: trees.splay.tests
+
+: randomize-numeric-splay-tree ( splay-tree -- )
+ 100 [ drop 100 random swap at drop ] with each ;
+
+: make-numeric-splay-tree ( n -- splay-tree )
+ <splay> [ [ conjoin ] curry each ] keep ;
+
+[ t ] [
+ 100 make-numeric-splay-tree dup randomize-numeric-splay-tree
+ [ [ drop , ] assoc-each ] { } make [ < ] monotonic?
+] unit-test
+
+[ 10 ] [ 10 make-numeric-splay-tree keys length ] unit-test
+[ 10 ] [ 10 make-numeric-splay-tree values length ] unit-test
+
+[ f ] [ <splay> f 4 pick set-at 4 swap at ] unit-test
+
+! Ensure that f can be a value
+[ t ] [ <splay> f 4 pick set-at 4 swap key? ] unit-test
+
+[
+{ { 1 "a" } { 2 "b" } { 3 "c" } { 4 "d" } { 5 "e" } { 6 "f" } }
+] [
+{
+ { 4 "d" } { 5 "e" } { 6 "f" }
+ { 1 "a" } { 2 "b" } { 3 "c" }
+} >splay >alist
+] unit-test
--- /dev/null
+! Copyright (c) 2005 Mackenzie Straight.
+! See http://factorcode.org/license.txt for BSD license.
+USING: arrays kernel math namespaces sequences assocs parser
+prettyprint.backend trees generic math.order accessors ;
+IN: trees.splay
+
+TUPLE: splay < tree ;
+
+: <splay> ( -- tree )
+ \ splay new-tree ;
+
+: rotate-right ( node -- node )
+ dup left>>
+ [ right>> swap (>>left) ] 2keep
+ [ (>>right) ] keep ;
+
+: rotate-left ( node -- node )
+ dup right>>
+ [ left>> swap (>>right) ] 2keep
+ [ (>>left) ] keep ;
+
+: link-right ( left right key node -- left right key node )
+ swap [ [ swap (>>left) ] 2keep
+ nip dup left>> ] dip swap ;
+
+: link-left ( left right key node -- left right key node )
+ swap [ rot [ (>>right) ] 2keep
+ drop dup right>> swapd ] dip swap ;
+
+: cmp ( key node -- obj node -1/0/1 )
+ 2dup key>> key-side ;
+
+: lcmp ( key node -- obj node -1/0/1 )
+ 2dup left>> key>> key-side ;
+
+: rcmp ( key node -- obj node -1/0/1 )
+ 2dup right>> key>> key-side ;
+
+DEFER: (splay)
+
+: splay-left ( left right key node -- left right key node )
+ dup left>> [
+ lcmp 0 < [ rotate-right ] when
+ dup left>> [ link-right (splay) ] when
+ ] when ;
+
+: splay-right ( left right key node -- left right key node )
+ dup right>> [
+ rcmp 0 > [ rotate-left ] when
+ dup right>> [ link-left (splay) ] when
+ ] when ;
+
+: (splay) ( left right key node -- left right key node )
+ cmp dup 0 <
+ [ drop splay-left ] [ 0 > [ splay-right ] when ] if ;
+
+: assemble ( head left right node -- root )
+ [ right>> swap (>>left) ] keep
+ [ left>> swap (>>right) ] keep
+ [ swap left>> swap (>>right) ] 2keep
+ [ swap right>> swap (>>left) ] keep ;
+
+: splay-at ( key node -- node )
+ [ T{ node } clone dup dup ] 2dip
+ (splay) nip assemble ;
+
+: splay ( key tree -- )
+ [ root>> splay-at ] keep (>>root) ;
+
+: splay-split ( key tree -- node node )
+ 2dup splay root>> cmp 0 < [
+ nip dup left>> swap f over (>>left)
+ ] [
+ nip dup right>> swap f over (>>right) swap
+ ] if ;
+
+: get-splay ( key tree -- node ? )
+ 2dup splay root>> cmp 0 = [
+ nip t
+ ] [
+ 2drop f f
+ ] if ;
+
+: get-largest ( node -- node )
+ dup [ dup right>> [ nip get-largest ] when* ] when ;
+
+: splay-largest ( node -- node )
+ dup [ dup get-largest key>> swap splay-at ] when ;
+
+: splay-join ( n2 n1 -- node )
+ splay-largest [
+ [ (>>right) ] keep
+ ] [
+ drop f
+ ] if* ;
+
+: remove-splay ( key tree -- )
+ tuck get-splay nip [
+ dup dec-count
+ dup right>> swap left>> splay-join
+ swap (>>root)
+ ] [ drop ] if* ;
+
+: set-splay ( value key tree -- )
+ 2dup get-splay [ 2nip (>>value) ] [
+ drop dup inc-count
+ 2dup splay-split rot
+ [ [ swapd ] dip node boa ] dip (>>root)
+ ] if ;
+
+: new-root ( value key tree -- )
+ 1 >>count
+ [ swap <node> ] dip (>>root) ;
+
+M: splay set-at ( value key tree -- )
+ dup root>> [ set-splay ] [ new-root ] if ;
+
+M: splay at* ( key tree -- value ? )
+ dup root>> [
+ get-splay [ dup [ value>> ] when ] dip
+ ] [
+ 2drop f f
+ ] if ;
+
+M: splay delete-at ( key tree -- )
+ dup root>> [ remove-splay ] [ 2drop ] if ;
+
+M: splay new-assoc
+ 2drop <splay> ;
+
+: >splay ( assoc -- tree )
+ T{ splay f f 0 } assoc-clone-like ;
+
+: SPLAY{
+ \ } [ >splay ] parse-literal ; parsing
+
+M: splay assoc-like
+ drop dup splay? [ >splay ] unless ;
+
+! M: splay pprint-delims drop \ SPLAY{ \ } ;
--- /dev/null
+Splay trees
--- /dev/null
+collections
+trees
--- /dev/null
+Binary search trees
--- /dev/null
+collections
+trees
--- /dev/null
+USING: help.syntax help.markup assocs ;
+IN: trees
+
+HELP: TREE{
+{ $syntax "TREE{ { key value }... }" }
+{ $values { "key" "a key" } { "value" "a value" } }
+{ $description "Literal syntax for an unbalanced tree." } ;
+
+HELP: <tree>
+{ $values { "tree" tree } }
+{ $description "Creates an empty unbalanced binary tree" } ;
+
+HELP: >tree
+{ $values { "assoc" assoc } { "tree" tree } }
+{ $description "Converts any " { $link assoc } " into an unbalanced binary tree." } ;
+
+HELP: tree
+{ $class-description "This is the class for unbalanced binary search trees. It is not usually intended to be used directly but rather as a basis for other trees." } ;
+
+ARTICLE: "trees" "Binary search trees"
+"This is a library for unbalanced binary search trees. It is not intended to be used directly in most situations but rather as a base class for new trees, because performance can degrade to linear time storage/retrieval by the number of keys. These binary search trees conform to the assoc protocol."
+{ $subsection tree }
+{ $subsection <tree> }
+{ $subsection >tree }
+{ $subsection POSTPONE: TREE{ } ;
+
+ABOUT: "trees"
--- /dev/null
+USING: trees assocs tools.test kernel sequences ;
+IN: trees.tests
+
+: test-tree ( -- tree )
+ TREE{
+ { 7 "seven" }
+ { 9 "nine" }
+ { 4 "four" }
+ { 4 "replaced four" }
+ { 7 "replaced seven" }
+ } clone ;
+
+! test set-at, at, at*
+[ "seven" ] [ <tree> "seven" 7 pick set-at 7 swap at ] unit-test
+[ "seven" t ] [ <tree> "seven" 7 pick set-at 7 swap at* ] unit-test
+[ f f ] [ <tree> "seven" 7 pick set-at 8 swap at* ] unit-test
+[ "seven" ] [ <tree> "seven" 7 pick set-at 7 swap at ] unit-test
+[ "replacement" ] [ <tree> "seven" 7 pick set-at "replacement" 7 pick set-at 7 swap at ] unit-test
+[ "replaced four" ] [ test-tree 4 swap at ] unit-test
+[ "nine" ] [ test-tree 9 swap at ] unit-test
+
+! test delete-at
+[ f ] [ test-tree 9 over delete-at 9 swap at ] unit-test
+[ "replaced seven" ] [ test-tree 9 over delete-at 7 swap at ] unit-test
+[ "replaced four" ] [ test-tree 9 over delete-at 4 swap at ] unit-test
+[ "nine" "replaced four" ] [ test-tree 7 over delete-at 9 over at 4 rot at ] unit-test
+[ "nine" ] [ test-tree 7 over delete-at 4 over delete-at 9 swap at ] unit-test
--- /dev/null
+! Copyright (C) 2007 Alex Chapman
+! See http://factorcode.org/license.txt for BSD license.
+USING: kernel generic math sequences arrays io namespaces
+prettyprint.private kernel.private assocs random combinators
+parser prettyprint.backend math.order accessors deques make
+prettyprint.custom ;
+IN: trees
+
+TUPLE: tree root count ;
+
+: new-tree ( class -- tree )
+ new
+ f >>root
+ 0 >>count ; inline
+
+: <tree> ( -- tree )
+ tree new-tree ;
+
+INSTANCE: tree assoc
+
+TUPLE: node key value left right ;
+
+: new-node ( key value class -- node )
+ new swap >>value swap >>key ;
+
+: <node> ( key value -- node )
+ node new-node ;
+
+SYMBOL: current-side
+
+: left ( -- symbol ) -1 ; inline
+: right ( -- symbol ) 1 ; inline
+
+: key-side ( k1 k2 -- n )
+ <=> {
+ { +lt+ [ -1 ] }
+ { +eq+ [ 0 ] }
+ { +gt+ [ 1 ] }
+ } case ;
+
+: go-left? ( -- ? ) current-side get left eq? ;
+
+: inc-count ( tree -- ) [ 1+ ] change-count drop ;
+
+: dec-count ( tree -- ) [ 1- ] change-count drop ;
+
+: node-link@ ( node ? -- node )
+ go-left? xor [ left>> ] [ right>> ] if ;
+: set-node-link@ ( left parent ? -- )
+ go-left? xor [ (>>left) ] [ (>>right) ] if ;
+
+: node-link ( node -- child ) f node-link@ ;
+: set-node-link ( child node -- ) f set-node-link@ ;
+: node+link ( node -- child ) t node-link@ ;
+: set-node+link ( child node -- ) t set-node-link@ ;
+
+: with-side ( side quot -- ) [ swap current-side set call ] with-scope ; inline
+: with-other-side ( quot -- )
+ current-side get neg swap with-side ; inline
+: go-left ( quot -- ) left swap with-side ; inline
+: go-right ( quot -- ) right swap with-side ; inline
+
+: leaf? ( node -- ? )
+ [ left>> ] [ right>> ] bi or not ;
+
+: random-side ( -- side ) left right 2array random ;
+
+: choose-branch ( key node -- key node-left/right )
+ 2dup key>> key-side [ node-link ] with-side ;
+
+: node-at* ( key node -- value ? )
+ [
+ 2dup key>> = [
+ nip value>> t
+ ] [
+ choose-branch node-at*
+ ] if
+ ] [ drop f f ] if* ;
+
+M: tree at* ( key tree -- value ? )
+ root>> node-at* ;
+
+: node-set ( value key node -- node )
+ 2dup key>> key-side dup 0 eq? [
+ drop nip swap >>value
+ ] [
+ [
+ [ node-link [ node-set ] [ swap <node> ] if* ] keep
+ [ set-node-link ] keep
+ ] with-side
+ ] if ;
+
+M: tree set-at ( value key tree -- )
+ [ [ node-set ] [ swap <node> ] if* ] change-root drop ;
+
+: valid-node? ( node -- ? )
+ [
+ dup dup left>> [ key>> swap key>> before? ] when*
+ [
+ dup dup right>> [ key>> swap key>> after? ] when* ] dip and swap
+ dup left>> valid-node? swap right>> valid-node? and and
+ ] [ t ] if* ;
+
+: valid-tree? ( tree -- ? ) root>> valid-node? ;
+
+: (node>alist) ( node -- )
+ [
+ [ left>> (node>alist) ]
+ [ [ key>> ] [ value>> ] bi 2array , ]
+ [ right>> (node>alist) ]
+ tri
+ ] when* ;
+
+M: tree >alist [ root>> (node>alist) ] { } make ;
+
+M: tree clear-assoc
+ 0 >>count
+ f >>root drop ;
+
+: copy-node-contents ( new old -- new )
+ [ key>> >>key ]
+ [ value>> >>value ] bi ;
+
+! Deletion
+DEFER: delete-node
+
+: (prune-extremity) ( parent node -- new-extremity )
+ dup node-link [
+ rot drop (prune-extremity)
+ ] [
+ tuck delete-node swap set-node-link
+ ] if* ;
+
+: prune-extremity ( node -- new-extremity )
+ #! remove and return the leftmost or rightmost child of this node.
+ #! assumes at least one child
+ dup node-link (prune-extremity) ;
+
+: replace-with-child ( node -- node )
+ dup node-link copy-node-contents dup node-link delete-node over set-node-link ;
+
+: replace-with-extremity ( node -- node )
+ dup node-link dup node+link [
+ ! predecessor/successor is not the immediate child
+ [ prune-extremity ] with-other-side copy-node-contents
+ ] [
+ ! node-link is the predecessor/successor
+ drop replace-with-child
+ ] if ;
+
+: delete-node-with-two-children ( node -- node )
+ #! randomised to minimise tree unbalancing
+ random-side [ replace-with-extremity ] with-side ;
+
+: delete-node ( node -- node )
+ #! delete this node, returning its replacement
+ dup left>> [
+ dup right>> [
+ delete-node-with-two-children
+ ] [
+ left>> ! left but no right
+ ] if
+ ] [
+ dup right>> [
+ right>> ! right but not left
+ ] [
+ drop f ! no children
+ ] if
+ ] if ;
+
+: delete-bst-node ( key node -- node )
+ 2dup key>> key-side dup 0 eq? [
+ drop nip delete-node
+ ] [
+ [ tuck node-link delete-bst-node over set-node-link ] with-side
+ ] if ;
+
+M: tree delete-at
+ [ delete-bst-node ] change-root drop ;
+
+M: tree new-assoc
+ 2drop <tree> ;
+
+M: tree clone dup assoc-clone-like ;
+
+: >tree ( assoc -- tree )
+ T{ tree f f 0 } assoc-clone-like ;
+
+M: tree assoc-like drop dup tree? [ >tree ] unless ;
+
+: TREE{
+ \ } [ >tree ] parse-literal ; parsing
+
+M: tree assoc-size count>> ;
+! M: tree pprint-delims drop \ TREE{ \ } ;
+! M: tree >pprint-sequence >alist ;
+! M: tree pprint-narrow? drop t ;
projects / factor.git / commitdiff