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 ;
8 TUPLE: tree root count ;
12 : new-tree ( class -- tree )
26 TUPLE: node key value left right ;
28 : new-node ( key value class -- node )
33 : <node> ( key value -- node )
41 : key-side ( k1 k2 -- n )
48 : go-left? ( -- ? ) current-side get left eq? ;
50 : inc-count ( tree -- ) [ 1 + ] change-count drop ;
52 : dec-count ( tree -- ) [ 1 - ] change-count drop ;
54 : node-link@ ( node ? -- node )
55 go-left? xor [ left>> ] [ right>> ] if ;
57 : set-node-link@ ( left parent ? -- )
58 go-left? xor [ (>>left) ] [ (>>right) ] if ;
60 : node-link ( node -- child ) f node-link@ ;
62 : set-node-link ( child node -- ) f set-node-link@ ;
64 : node+link ( node -- child ) t node-link@ ;
66 : set-node+link ( child node -- ) t set-node-link@ ;
68 : with-side ( side quot -- )
69 [ current-side ] dip with-variable ; inline
71 : with-other-side ( quot -- )
72 current-side get neg swap with-side ; inline
74 : go-left ( quot -- ) left swap with-side ; inline
76 : go-right ( quot -- ) right swap with-side ; inline
79 [ left>> ] [ right>> ] bi or not ;
81 : random-side ( -- side )
82 left right 2array random ;
84 : choose-branch ( key node -- key node-left/right )
85 2dup key>> key-side [ node-link ] with-side ;
87 : node-at* ( key node -- value ? )
92 choose-branch node-at*
96 M: tree at* ( key tree -- value ? )
99 : node-set ( value key node -- node )
100 2dup key>> key-side dup 0 eq? [
101 drop nip swap >>value
104 [ node-link [ node-set ] [ swap <node> ] if* ] keep
105 [ set-node-link ] keep
109 M: tree set-at ( value key tree -- )
110 [ [ node-set ] [ swap <node> ] if* ] change-root drop ;
112 : valid-node? ( node -- ? )
114 dup dup left>> [ key>> swap key>> before? ] when*
116 dup dup right>> [ key>> swap key>> after? ] when* ] dip and swap
117 dup left>> valid-node? swap right>> valid-node? and and
120 : valid-tree? ( tree -- ? ) root>> valid-node? ;
122 : (node>alist) ( node -- )
124 [ left>> (node>alist) ]
125 [ [ key>> ] [ value>> ] bi 2array , ]
126 [ right>> (node>alist) ]
130 M: tree >alist [ root>> (node>alist) ] { } make ;
136 : copy-node-contents ( new old -- new )
138 [ value>> >>value ] bi ;
143 : (prune-extremity) ( parent node -- new-extremity )
145 [ nip ] dip (prune-extremity)
147 [ delete-node ] [ set-node-link ] bi
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) ;
155 : replace-with-child ( node -- node )
156 dup node-link copy-node-contents dup node-link delete-node over set-node-link ;
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
163 ! node-link is the predecessor/successor
164 drop replace-with-child
167 : delete-node-with-two-children ( node -- node )
168 #! randomised to minimise tree unbalancing
169 random-side [ replace-with-extremity ] with-side ;
171 : delete-node ( node -- node )
172 #! delete this node, returning its replacement
175 delete-node-with-two-children
177 left>> ! left but no right
181 right>> ! right but not left
187 : delete-bst-node ( key node -- node )
188 2dup key>> key-side dup 0 eq? [
192 [ node-link delete-bst-node ]
201 [ delete-bst-node ] change-root drop ;
206 M: tree clone dup assoc-clone-like ;
208 : >tree ( assoc -- tree )
209 T{ tree f f 0 } assoc-clone-like ;
211 M: tree assoc-like drop dup tree? [ >tree ] unless ;
214 \ } [ >tree ] parse-literal ;
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 ;