459 lines
12 KiB
Factor
459 lines
12 KiB
Factor
! Copyright (C) 2007 Alex Chapman
|
|
! See http://factorcode.org/license.txt for BSD license.
|
|
USING: accessors arrays assocs combinators
|
|
combinators.short-circuit deques dlists kernel locals make math
|
|
math.order namespaces parser prettyprint.custom random sequences
|
|
vectors ;
|
|
IN: trees
|
|
|
|
TUPLE: tree root { count integer } ;
|
|
|
|
<PRIVATE
|
|
|
|
: new-tree ( class -- tree )
|
|
new
|
|
f >>root
|
|
0 >>count ; inline
|
|
|
|
PRIVATE>
|
|
|
|
: <tree> ( -- tree )
|
|
tree new-tree ;
|
|
|
|
INSTANCE: tree assoc
|
|
|
|
<PRIVATE
|
|
|
|
TUPLE: node key value left right ;
|
|
|
|
: new-node ( key value class -- node )
|
|
new
|
|
swap >>value
|
|
swap >>key ; inline
|
|
|
|
: <node> ( key value -- node )
|
|
node new-node ;
|
|
|
|
SYMBOL: current-side
|
|
|
|
CONSTANT: left -1
|
|
CONSTANT: right 1
|
|
|
|
: key-side ( k1 k2 -- n )
|
|
<=> {
|
|
{ +lt+ [ left ] }
|
|
{ +eq+ [ 0 ] }
|
|
{ +gt+ [ right ] }
|
|
} 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 -- )
|
|
[ current-side ] dip with-variable ; 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>> not ] [ right>> not ] } 1&& ;
|
|
|
|
: random-side ( -- side )
|
|
2 random 0 eq? left right ? ;
|
|
|
|
: 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*
|
|
root>> node-at* ;
|
|
|
|
: node-set ( value key node -- node new? )
|
|
2dup key>> key-side dup 0 eq? [
|
|
drop nip swap >>value f
|
|
] [
|
|
[
|
|
[ node-link [ node-set ] [ swap <node> t ] if* ] keep
|
|
swap [ [ set-node-link ] keep ] dip
|
|
] with-side
|
|
] if ;
|
|
|
|
M: tree set-at
|
|
[ [ node-set ] [ swap <node> t ] if* swap ] change-root
|
|
swap [ dup inc-count ] when drop ;
|
|
|
|
: valid-node? ( node -- ? )
|
|
[
|
|
{
|
|
[ dup left>> [ key>> swap key>> before? ] when* ]
|
|
[ dup right>> [ key>> swap key>> after? ] when* ]
|
|
[ left>> valid-node? ]
|
|
[ right>> valid-node? ]
|
|
} 1&&
|
|
] [ t ] if* ;
|
|
|
|
: valid-tree? ( tree -- ? ) root>> valid-node? ;
|
|
|
|
: node>entry ( node -- entry ) [ key>> ] [ value>> ] bi 2array ;
|
|
|
|
: entry, ( node -- ) node>entry , ;
|
|
|
|
: (node>alist) ( node -- )
|
|
[
|
|
[ left>> (node>alist) ]
|
|
[ entry, ]
|
|
[ right>> (node>alist) ]
|
|
tri
|
|
] when* ;
|
|
|
|
M: tree >alist
|
|
[ root>> (node>alist) ] { } make ;
|
|
|
|
:: (node>subalist-right) ( to-key node end-comparator: ( key1 key2 -- ? ) -- )
|
|
node [
|
|
node key>> to-key end-comparator call :> node-left?
|
|
|
|
node left>> node-left? [ (node>alist) ] [
|
|
[ to-key ] dip end-comparator (node>subalist-right)
|
|
] if
|
|
|
|
node-left? [
|
|
node [ entry, ] [
|
|
right>> [ to-key ] dip
|
|
end-comparator (node>subalist-right)
|
|
] bi
|
|
] when
|
|
] when ; inline recursive
|
|
|
|
:: (node>subalist-left) ( from-key node start-comparator: ( key1 key2 -- ? ) -- )
|
|
node [
|
|
node key>> from-key start-comparator call :> node-right?
|
|
|
|
node-right? [
|
|
node [
|
|
left>> [ from-key ] dip
|
|
start-comparator (node>subalist-left)
|
|
] [ entry, ] bi
|
|
] when
|
|
|
|
node right>> node-right? [ (node>alist) ] [
|
|
[ from-key ] dip start-comparator (node>subalist-left)
|
|
] if
|
|
] when ; inline recursive
|
|
|
|
:: (node>subalist) ( from-key to-key node start-comparator: ( key1 key2 -- ? ) end-comparator: ( key1 key2 -- ? ) -- )
|
|
node [
|
|
node key>> from-key start-comparator call :> node-right?
|
|
node key>> to-key end-comparator call :> node-left?
|
|
|
|
node-right? [
|
|
from-key node left>> node-left?
|
|
[ start-comparator (node>subalist-left) ]
|
|
[
|
|
[ to-key ] dip start-comparator
|
|
end-comparator (node>subalist)
|
|
] if
|
|
] when
|
|
|
|
node-right? node-left? and [ node entry, ] when
|
|
|
|
node-left? [
|
|
to-key node right>> node-right?
|
|
[ end-comparator (node>subalist-right) ]
|
|
[
|
|
[ from-key ] 2dip start-comparator
|
|
end-comparator (node>subalist)
|
|
] if
|
|
] when
|
|
] when ; inline recursive
|
|
|
|
PRIVATE>
|
|
|
|
: subtree>alist[) ( from-key to-key tree -- alist )
|
|
[ root>> [ after=? ] [ before? ] (node>subalist) ] { } make ;
|
|
|
|
: subtree>alist(] ( from-key to-key tree -- alist )
|
|
[ root>> [ after? ] [ before=? ] (node>subalist) ] { } make ;
|
|
|
|
: subtree>alist[] ( from-key to-key tree -- alist )
|
|
[ root>> [ after=? ] [ before=? ] (node>subalist) ] { } make ;
|
|
|
|
: subtree>alist() ( from-key to-key tree -- alist )
|
|
[ root>> [ after? ] [ before? ] (node>subalist) ] { } make ;
|
|
|
|
: headtree>alist[) ( to-key tree -- alist )
|
|
[ root>> [ before? ] (node>subalist-right) ] { } make ;
|
|
|
|
: headtree>alist[] ( to-key tree -- alist )
|
|
[ root>> [ before=? ] (node>subalist-right) ] { } make ;
|
|
|
|
: tailtree>alist[] ( from-key tree -- alist )
|
|
[ root>> [ after=? ] (node>subalist-left) ] { } make ;
|
|
|
|
: tailtree>alist(] ( from-key tree -- alist )
|
|
[ root>> [ after? ] (node>subalist-left) ] { } make ;
|
|
|
|
<PRIVATE
|
|
|
|
: (nodepath-at) ( key node -- )
|
|
[
|
|
dup ,
|
|
2dup key>> = [
|
|
2drop
|
|
] [
|
|
choose-branch (nodepath-at)
|
|
] if
|
|
] [ drop ] if* ;
|
|
|
|
: nodepath-at ( key tree -- path )
|
|
[ root>> (nodepath-at) ] { } make ;
|
|
|
|
: right-extremity ( node -- node' )
|
|
[ dup right>> dup ] [ nip ] while drop ;
|
|
|
|
: left-extremity ( node -- node' )
|
|
[ dup left>> dup ] [ nip ] while drop ;
|
|
|
|
: lower-node-in-child? ( key node -- ? )
|
|
[ nip left>> ] [ key>> = ] 2bi and ;
|
|
|
|
: higher-node-in-child? ( key node -- ? )
|
|
[ nip right>> ] [ key>> = ] 2bi and ;
|
|
|
|
: lower-node ( key tree -- node )
|
|
dupd nodepath-at
|
|
[ drop f ] [
|
|
reverse 2dup first lower-node-in-child?
|
|
[ nip first left>> right-extremity ]
|
|
[ [ key>> after? ] with find nip ] if
|
|
] if-empty ;
|
|
|
|
: higher-node ( key tree -- node )
|
|
dupd nodepath-at
|
|
[ drop f ] [
|
|
reverse 2dup first higher-node-in-child?
|
|
[ nip first right>> left-extremity ]
|
|
[ [ key>> before? ] with find nip ] if
|
|
] if-empty ;
|
|
|
|
: floor-node ( key tree -- node )
|
|
dupd nodepath-at [ drop f ] [
|
|
reverse [ key>> after=? ] with find nip
|
|
] if-empty ;
|
|
|
|
: ceiling-node ( key tree -- node )
|
|
dupd nodepath-at [ drop f ] [
|
|
reverse [ key>> before=? ] with find nip
|
|
] if-empty ;
|
|
|
|
: first-node ( tree -- node ) root>> dup [ left-extremity ] when ;
|
|
|
|
: last-node ( tree -- node ) root>> dup [ right-extremity ] when ;
|
|
|
|
PRIVATE>
|
|
|
|
: lower-entry ( key tree -- pair/f ) lower-node dup [ node>entry ] when ;
|
|
|
|
: higher-entry ( key tree -- pair/f ) higher-node dup [ node>entry ] when ;
|
|
|
|
: floor-entry ( key tree -- pair/f ) floor-node dup [ node>entry ] when ;
|
|
|
|
: ceiling-entry ( key tree -- pair/f ) ceiling-node dup [ node>entry ] when ;
|
|
|
|
: first-entry ( tree -- pair/f ) first-node dup [ node>entry ] when ;
|
|
|
|
: last-entry ( tree -- pair/f ) last-node dup [ node>entry ] when ;
|
|
|
|
: lower-key ( key tree -- key/f ) lower-node dup [ key>> ] when ;
|
|
|
|
: higher-key ( key tree -- key/f ) higher-node dup [ key>> ] when ;
|
|
|
|
: floor-key ( key tree -- key/f ) floor-node dup [ key>> ] when ;
|
|
|
|
: ceiling-key ( key tree -- key/f ) ceiling-node dup [ key>> ] when ;
|
|
|
|
: first-key ( tree -- key/f ) first-node dup [ key>> ] when ;
|
|
|
|
: last-key ( tree -- key/f ) last-node dup [ key>> ] when ;
|
|
|
|
<PRIVATE
|
|
|
|
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 [
|
|
nipd (prune-extremity)
|
|
] [
|
|
[ delete-node swap set-node-link ] keep
|
|
] 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 [ right>> ] [ left>> ] bi [
|
|
swap [
|
|
drop delete-node-with-two-children
|
|
] [
|
|
nip ! left but no right
|
|
] if
|
|
] [
|
|
nip ! right but no left, or no children
|
|
] if* ;
|
|
|
|
: delete-bst-node ( key node -- node deleted? )
|
|
2dup key>> key-side dup 0 eq? [
|
|
drop nip delete-node t
|
|
] [
|
|
[
|
|
[ node-link delete-bst-node ]
|
|
[ swap [ set-node-link ] dip ]
|
|
[ swap ] tri
|
|
] with-side
|
|
] if ;
|
|
|
|
M: tree delete-at
|
|
[ delete-bst-node swap ] change-root
|
|
swap [ dup dec-count ] when drop ;
|
|
|
|
M: tree new-assoc
|
|
2drop <tree> ;
|
|
|
|
: clone-nodes ( node -- node' )
|
|
dup [
|
|
clone
|
|
[ clone-nodes ] change-left
|
|
[ clone-nodes ] change-right
|
|
] when ;
|
|
|
|
M: tree clone (clone) [ clone-nodes ] change-root ;
|
|
|
|
: ?push-children ( node queue -- )
|
|
[ [ left>> ] [ right>> ] bi ]
|
|
[ [ over [ push-front ] [ 2drop ] if ] curry bi@ ] bi* ;
|
|
|
|
: each-bfs-node ( tree quot: ( ... entry -- ... ) -- ... )
|
|
[ root>> <dlist> [ push-front ] keep dup ] dip
|
|
[
|
|
[ drop node>entry ] prepose
|
|
[ ?push-children ] 2bi
|
|
] 2curry slurp-deque ; inline
|
|
|
|
: >bfs-alist ( tree -- alist )
|
|
dup assoc-size <vector> [
|
|
[ push ] curry each-bfs-node
|
|
] keep ;
|
|
|
|
M: tree assoc-clone-like
|
|
[ dup tree? [ >bfs-alist ] when ] dip call-next-method ;
|
|
|
|
PRIVATE>
|
|
|
|
: >tree ( assoc -- tree )
|
|
T{ tree f f 0 } assoc-clone-like ;
|
|
|
|
SYNTAX: TREE{
|
|
\ } [ >tree ] parse-literal ;
|
|
|
|
<PRIVATE
|
|
|
|
M: tree assoc-like drop dup tree? [ >tree ] unless ;
|
|
|
|
M: tree assoc-size count>> ;
|
|
M: tree pprint-delims drop \ TREE{ \ } ;
|
|
M: tree >pprint-sequence >alist ;
|
|
M: tree pprint-narrow? drop t ;
|
|
|
|
: node-height ( node -- n )
|
|
[
|
|
[ left>> ] [ right>> ] bi
|
|
[ node-height ] bi@ max 1 +
|
|
] [ 0 ] if* ;
|
|
|
|
PRIVATE>
|
|
|
|
: height ( tree -- n )
|
|
root>> node-height ;
|
|
|
|
<PRIVATE
|
|
|
|
: pop-tree-extremity ( tree node/f -- node/f )
|
|
dup [
|
|
[ key>> swap delete-at ] keep node>entry
|
|
] [ nip ] if ;
|
|
|
|
:: slurp-tree ( tree quot: ( ... entry -- ... ) getter: ( tree -- node ) -- ... )
|
|
[ tree count>> 0 = ]
|
|
[ tree getter call quot call ] until ; inline
|
|
|
|
PRIVATE>
|
|
|
|
: pop-tree-left ( tree -- node/f )
|
|
dup first-node pop-tree-extremity ;
|
|
|
|
: pop-tree-right ( tree -- node/f )
|
|
dup last-node pop-tree-extremity ;
|
|
|
|
: slurp-tree-left ( tree quot: ( ... entry -- ... ) -- ... )
|
|
[ pop-tree-left ] slurp-tree ; inline
|
|
|
|
: slurp-tree-right ( tree quot: ( ... entry -- ... ) -- ... )
|
|
[ pop-tree-right ] slurp-tree ; inline
|