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