Paste: a binary-tree implementation in factor

USING: accessors fry kernel locals make math multiline
prettyprint sequences shuffle slots.syntax ;
QUALIFIED: math
! QUALIFIED-WITH: math.order mo
IN: binary-tree

TUPLE: node parent left right { size-below initial: 0 } value ;
TUPLE: binary-tree { root node } ;


! : <=> ( node1 node2 -- ? )
!    [ value>> ] bi@ mo:<=> ;
! : > ( node1 node2 -- ? ) 
!    [ value>> ] bi@ math:> ; 

! : = ( node1 node2 -- ? )
!     [ value>> ] bi@ kernel:= ;

: < ( node1 node2 -- ? ) 
    [ value>> ] bi@ math:< ;
    
: >= ( node1 node2 -- ? ) 
    < not ;
! : <= ( node1 node2 -- ? ) 
!    > not ;


: <node> ( value -- node )    
    node new 
    swap >>value ;
    
: <binary-tree> ( node -- btree )
    binary-tree new swap >>root ;

    
<PRIVATE

: inc-parent-size ( child -- )
    parent>> dup 
    [
        [ 1 + ] change-size-below inc-parent-size 
    ] [ drop ] if ;
    
: add-child-on-right ( node child -- )
    over >>parent 
        [ >>right drop ]
        [ nip inc-parent-size ] 2bi ;
: add-child-on-left ( node child -- )
    over >>parent 
        [ >>left drop ]
        [ nip inc-parent-size ] 2bi ;
PRIVATE>

GENERIC: add-node ( child node -- )


M: node add-node ( child node -- )
    2dup >= [ ! go right
        dup right>> ! can be node/f
            [ right>> add-node ] [ swap add-child-on-right ] if
    ] [ ! go left
        dup left>> 
            [ left>> add-node ] [ swap add-child-on-left ] if
    ] if ;

M: binary-tree add-node ( child node -- )
    root>> add-node ;
    
: add-value ( node value -- )
    <node> swap add-node ;
    
GENERIC: min ( node -- min-node )
GENERIC: max ( node -- max-node )

M: node min ( node -- min-node )
    dup left>> [ left>> min ] when ;

M: node max ( node -- max-node )
    dup right>> [ right>> max ] when ;

M: binary-tree min ( node -- min-node )
    root>> min ;
    
M: binary-tree max ( node -- max-node )
    root>> max ;

: root ( node -- root )
    dup parent>> [ parent>> root ] when ;
    
<PRIVATE
: (get-list) ( node -- )
    [ left>> dup [ (get-list) ] [ drop ] if ]
    [ , ] 
    [ right>> dup [ (get-list) ] [ drop ] if ] tri ;
PRIVATE>   

GENERIC: items ( base-node -- items  )
M: node items ( base-node -- seq )
    [ (get-list) ] { } make ;
M: binary-tree items ( base-node -- seq )    
    root>> items ;

: values ( btree -- seq )
    items [ value>> ] map ;
    
: print ( binary-tree -- )
    values . ;
    
GENERIC: find-node* ( value node -- node/f )    

: find-node ( value node -- node/f )    
    dup [ find-node* ] [ 2drop f ] if ;
    
M:: node find-node* ( value node -- node/f )    
    node value>> value = 
    [ node ]
    [ value node left>> find-node :> left-result
        left-result dup [ drop value node right>> find-node :> right-result 
        right-result dup [ drop f ] unless
        ] unless
    ] if ;
    
M: binary-tree find-node* ( value btree -- node/f )
    root>> find-node ;    

    
: add-values ( node seq -- node' )
    dupd [ add-value ] with each ;

: add-nodes ( items node -- node' )
    tuck '[ _ add-node ] each ;
    
: children ( node -- seq )
    slots{ right left } ;

: right? ( child -- ? )
    dup parent>> right>> eq? ;

<PRIVATE
: rm-child ( node -- )
    [ parent>> ] [ right? ] bi 
    [ f >>right ] [ f >>left ] if drop ;
PRIVATE>

: remove ( node -- )
    dup [
    [ children harvest ] [ parent>> ] [ rm-child ] tri add-nodes 
    ] when drop ;

: remove-nodes ( nodes -- )
    [ remove ] each ;
    
: find-nodes ( btree values -- nodes )
    swap [ find-node ] curry map ;
    
: remove-values ( btree values -- )
    find-nodes remove-nodes ;

Annotation: can't remove root node

the above code doesn't handle the case where the user wants to remove the root node.
right now, it asks the node's parent to add node's children.
but a root node doesn't have a parent. it would be f.

I read that on balanced trees, they take the left child's max, add its children (left or none) to its parent, then back to the top of the tree, use that node as the new root, and add the previous left and right to it.
it's pretty good, because the order is kept.

here's untested code

: root-remove ( node -- )
    [ children ] [ left max dup remove f >>right f >>left ] bi add-nodes ;

: normal-remove ( node -- )
    dup [
    [ children harvest ] [ parent>> ] [ rm-child ] tri add-nodes 
    ] when drop ;
    
: remove ( node -- )
    dup root eq?
    [ root-remove ] [ normal-remove ] if ;

New Annotation