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