extra/trees/trees.factor

   1 ! Copyright (C) 2007 Alex Chapman
   2 ! See http://factorcode.org/license.txt for BSD license.
   3 USING: kernel generic math sequences arrays io namespaces
   4 prettyprint.private kernel.private assocs random combinators
   5 parser prettyprint.backend math.order accessors ;
   6 IN: trees
   7
   8 TUPLE: tree root count ;
   9
  10 : new-tree ( class -- tree )
  11     new
  12         f >>root
  13         0 >>count ; inline
  14
  15 : <tree> ( -- tree )
  16     tree new-tree ;
  17
  18 INSTANCE: tree assoc
  19
  20 TUPLE: node key value left right ;
  21
  22 : new-node ( key value class -- node )
  23     new swap >>value swap >>key ;
  24
  25 : <node> ( key value -- node )
  26     node new-node ;
  27
  28 SYMBOL: current-side
  29
  30 : left ( -- symbol ) -1 ; inline
  31 : right ( -- symbol ) 1 ; inline
  32
  33 : key-side ( k1 k2 -- n )
  34     <=> {
  35         { +lt+ [ -1 ] }
  36         { +eq+ [ 0 ] }
  37         { +gt+ [ 1 ] }
  38     } case ;
  39
  40 : go-left? ( -- ? ) current-side get left eq? ;
  41
  42 : inc-count ( tree -- ) [ 1+ ] change-count drop ;
  43
  44 : dec-count ( tree -- ) [ 1- ] change-count drop ;
  45
  46 : node-link@ ( node ? -- node )
  47     go-left? xor [ left>> ] [ right>> ] if ;
  48 : set-node-link@ ( left parent ? -- )
  49     go-left? xor [ set-node-left ] [ set-node-right ] if ;
  50
  51 : node-link ( node -- child ) f node-link@  ;
  52 : set-node-link ( child node -- ) f set-node-link@ ;
  53 : node+link ( node -- child ) t node-link@ ;
  54 : set-node+link ( child node -- ) t set-node-link@ ;
  55
  56 : with-side ( side quot -- ) [ swap current-side set call ] with-scope ; inline
  57 : with-other-side ( quot -- )
  58     current-side get neg swap with-side ; inline
  59 : go-left ( quot -- ) left swap with-side ; inline
  60 : go-right ( quot -- ) right swap with-side ; inline
  61
  62 : leaf? ( node -- ? )
  63     [ left>> ] [ right>> ] bi or not ;
  64
  65 : random-side ( -- side ) left right 2array random ;
  66
  67 : choose-branch ( key node -- key node-left/right )
  68     2dup node-key key-side [ node-link ] with-side ;
  69
  70 : node-at* ( key node -- value ? )
  71     [
  72         2dup node-key = [
  73             nip node-value t
  74         ] [
  75             choose-branch node-at*
  76         ] if
  77     ] [ drop f f ] if* ;
  78
  79 M: tree at* ( key tree -- value ? )
  80     root>> node-at* ;
  81
  82 : node-set ( value key node -- node )
  83     2dup key>> key-side dup 0 eq? [
  84         drop nip swap >>value
  85     ] [
  86         [
  87             [ node-link [ node-set ] [ swap <node> ] if* ] keep
  88             [ set-node-link ] keep
  89         ] with-side
  90     ] if ;
  91
  92 M: tree set-at ( value key tree -- )
  93     [ [ node-set ] [ swap <node> ] if* ] change-root drop ;
  94
  95 : valid-node? ( node -- ? )
  96     [
  97         dup dup left>> [ node-key swap node-key before? ] when* >r
  98         dup dup right>> [ node-key swap node-key after? ] when* r> and swap
  99         dup left>> valid-node? swap right>> valid-node? and and
 100     ] [ t ] if* ;
 101
 102 : valid-tree? ( tree -- ? ) root>> valid-node? ;
 103
 104 : (node>alist) ( node -- )
 105     [
 106         [ left>> (node>alist) ]
 107         [ [ node-key ] [ node-value ] bi 2array , ]
 108         [ right>> (node>alist) ]
 109         tri
 110     ] when* ;
 111
 112 M: tree >alist [ root>> (node>alist) ] { } make ;
 113
 114 M: tree clear-assoc
 115     0 >>count
 116     f >>root drop ;
 117
 118 : copy-node-contents ( new old -- )
 119     dup node-key pick set-node-key node-value swap set-node-value ;
 120
 121 ! Deletion
 122 DEFER: delete-node
 123
 124 : (prune-extremity) ( parent node -- new-extremity )
 125     dup node-link [
 126         rot drop (prune-extremity)
 127     ] [
 128         tuck delete-node swap set-node-link
 129     ] if* ;
 130
 131 : prune-extremity ( node -- new-extremity )
 132     #! remove and return the leftmost or rightmost child of this node.
 133     #! assumes at least one child
 134     dup node-link (prune-extremity) ;
 135
 136 : replace-with-child ( node -- node )
 137     dup dup node-link copy-node-contents dup node-link delete-node over set-node-link ;
 138
 139 : replace-with-extremity ( node -- node )
 140     dup node-link dup node+link [
 141         ! predecessor/successor is not the immediate child
 142         [ prune-extremity ] with-other-side dupd copy-node-contents
 143     ] [
 144         ! node-link is the predecessor/successor
 145         drop replace-with-child
 146     ] if ;
 147
 148 : delete-node-with-two-children ( node -- node )
 149     #! randomised to minimise tree unbalancing
 150     random-side [ replace-with-extremity ] with-side ;
 151
 152 : delete-node ( node -- node )
 153     #! delete this node, returning its replacement
 154     dup left>> [
 155         dup right>> [
 156             delete-node-with-two-children
 157         ] [
 158             left>> ! left but no right
 159         ] if
 160     ] [
 161         dup right>> [
 162             right>> ! right but not left
 163         ] [
 164             drop f ! no children
 165         ] if
 166     ] if ;
 167
 168 : delete-bst-node ( key node -- node )
 169     2dup node-key key-side dup 0 eq? [
 170         drop nip delete-node
 171     ] [
 172         [ tuck node-link delete-bst-node over set-node-link ] with-side
 173     ] if ;
 174
 175 M: tree delete-at
 176     [ delete-bst-node ] change-root drop ;
 177
 178 M: tree new-assoc
 179     2drop <tree> ;
 180
 181 M: tree clone dup assoc-clone-like ;
 182
 183 : >tree ( assoc -- tree )
 184     T{ tree f f 0 } assoc-clone-like ;
 185
 186 M: tree assoc-like drop dup tree? [ >tree ] unless ;
 187
 188 : TREE{
 189     \ } [ >tree ] parse-literal ; parsing
 190
 191 M: tree pprint-delims drop \ TREE{ \ } ;
 192 M: tree assoc-size count>> ;
 193 M: tree >pprint-sequence >alist ;
 194 M: tree pprint-narrow? drop t ;