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