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