extra/trees/trees.factor

   1 ! Copyright (C) 2007 Alex Chapman
   2 ! See http://factorcode.org/license.txt for BSD license.
   3 USING: accessors arrays assocs combinators
   4 combinators.short-circuit kernel make math math.order namespaces
   5 parser prettyprint.custom random ;
   6 IN: trees
   7
   8 TUPLE: tree root { count integer } ;
   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 ; inline
  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+ [ left ] }
  44         { +eq+ [ 0 ] }
  45         { +gt+ [ right ] }
  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>> not ] [ right>> not ] } 1&& ;
  80
  81 : random-side ( -- side )
  82     2 random 0 eq? left right ? ;
  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*
  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
 110     [ [ node-set ] [ swap <node> ] if* ] change-root drop ;
 111
 112 : valid-node? ( node -- ? )
 113     [
 114         {
 115             [ dup left>> [ key>> swap key>> before? ] when* ]
 116             [ dup right>> [ key>> swap key>> after? ] when* ]
 117             [ left>> valid-node? ]
 118             [ right>> valid-node? ]
 119         } 1&&
 120     ] [ t ] if* ;
 121
 122 : valid-tree? ( tree -- ? ) root>> valid-node? ;
 123
 124 : (node>alist) ( node -- )
 125     [
 126         [ left>> (node>alist) ]
 127         [ [ key>> ] [ value>> ] bi 2array , ]
 128         [ right>> (node>alist) ]
 129         tri
 130     ] when* ;
 131
 132 M: tree >alist
 133     [ root>> (node>alist) ] { } make ;
 134
 135 M: tree clear-assoc
 136     0 >>count
 137     f >>root drop ;
 138
 139 : copy-node-contents ( new old -- new )
 140     [ key>> >>key ]
 141     [ value>> >>value ] bi ;
 142
 143 ! Deletion
 144 DEFER: delete-node
 145
 146 : (prune-extremity) ( parent node -- new-extremity )
 147     dup node-link [
 148         [ nip ] dip (prune-extremity)
 149     ] [
 150         [ delete-node ] [ set-node-link ] bi
 151     ] if* ;
 152
 153 : prune-extremity ( node -- new-extremity )
 154     ! remove and return the leftmost or rightmost child of this node.
 155     ! assumes at least one child
 156     dup node-link (prune-extremity) ;
 157
 158 : replace-with-child ( node -- node )
 159     dup node-link copy-node-contents dup node-link delete-node over set-node-link ;
 160
 161 : replace-with-extremity ( node -- node )
 162     dup node-link dup node+link [
 163         ! predecessor/successor is not the immediate child
 164         [ prune-extremity ] with-other-side copy-node-contents
 165     ] [
 166         ! node-link is the predecessor/successor
 167         drop replace-with-child
 168     ] if ;
 169
 170 : delete-node-with-two-children ( node -- node )
 171     ! randomised to minimise tree unbalancing
 172     random-side [ replace-with-extremity ] with-side ;
 173
 174 : delete-node ( node -- node )
 175     ! delete this node, returning its replacement
 176     dup [ right>> ] [ left>> ] bi [
 177         swap [
 178             drop delete-node-with-two-children
 179         ] [
 180             nip ! left but no right
 181         ] if
 182     ] [
 183         nip ! right but no left, or no children
 184     ] if* ;
 185
 186 : delete-bst-node ( key node -- node )
 187     2dup key>> key-side dup 0 eq? [
 188         drop nip delete-node
 189     ] [
 190         [
 191             [ node-link delete-bst-node ]
 192             [ set-node-link ]
 193             [ ] tri
 194         ] with-side
 195     ] if ;
 196
 197 PRIVATE>
 198
 199 M: tree delete-at
 200     [ delete-bst-node ] change-root drop ;
 201
 202 M: tree new-assoc
 203     2drop <tree> ;
 204
 205 M: tree clone dup assoc-clone-like ;
 206
 207 : >tree ( assoc -- tree )
 208     T{ tree f f 0 } assoc-clone-like ;
 209
 210 M: tree assoc-like drop dup tree? [ >tree ] unless ;
 211
 212 SYNTAX: TREE{
 213     \ } [ >tree ] parse-literal ;
 214
 215 M: tree assoc-size count>> ;
 216 M: tree pprint-delims drop \ TREE{ \ } ;
 217 M: tree >pprint-sequence >alist ;
 218 M: tree pprint-narrow? drop t ;