1 ! Copyright (C) 2007 Alex Chapman
2 ! See https://factorcode.org/license.txt for BSD license.
3 USING: accessors arrays assocs combinators
4 combinators.short-circuit deques dlists kernel make math
5 math.order namespaces parser prettyprint.custom random sequences
9 TUPLE: tree root { count integer } ;
13 : new-tree ( class -- tree )
27 TUPLE: node key value left right ;
29 : new-node ( key value class -- node )
34 : <node> ( key value -- node )
42 : key-side ( k1 k2 -- n )
49 : go-left? ( -- ? ) current-side get left eq? ;
51 : inc-count ( tree -- ) [ 1 + ] change-count drop ;
53 : dec-count ( tree -- ) [ 1 - ] change-count drop ;
55 : node-link@ ( node ? -- node )
56 go-left? xor [ left>> ] [ right>> ] if ;
58 : set-node-link@ ( left parent ? -- )
59 go-left? xor [ left<< ] [ right<< ] if ;
61 : node-link ( node -- child ) f node-link@ ;
63 : set-node-link ( child node -- ) f set-node-link@ ;
65 : node+link ( node -- child ) t node-link@ ;
67 : set-node+link ( child node -- ) t set-node-link@ ;
69 : with-side ( side quot -- )
70 [ current-side ] dip with-variable ; inline
72 : with-other-side ( quot -- )
73 current-side get neg swap with-side ; inline
75 : go-left ( quot -- ) left swap with-side ; inline
77 : go-right ( quot -- ) right swap with-side ; inline
80 { [ left>> not ] [ right>> not ] } 1&& ;
82 : random-side ( -- side )
83 2 random 0 eq? left right ? ;
85 : choose-branch ( key node -- key node-left/right )
86 2dup key>> key-side [ node-link ] with-side ;
88 : node-at* ( key node -- value ? )
93 choose-branch node-at*
100 : node-set ( value key node -- node new? )
101 2dup key>> key-side dup 0 eq? [
102 drop nip swap >>value f
105 [ node-link [ node-set ] [ swap <node> t ] if* ] keep
106 swap [ [ set-node-link ] keep ] dip
111 [ [ node-set ] [ swap <node> t ] if* swap ] change-root
112 swap [ dup inc-count ] when drop ;
114 : valid-node? ( node -- ? )
117 [ dup left>> [ key>> swap key>> before? ] when* ]
118 [ dup right>> [ key>> swap key>> after? ] when* ]
119 [ left>> valid-node? ]
120 [ right>> valid-node? ]
124 : valid-tree? ( tree -- ? ) root>> valid-node? ;
126 : node>entry ( node -- entry ) [ key>> ] [ value>> ] bi 2array ;
128 : entry, ( node -- ) node>entry , ;
130 : (node>alist) ( node -- )
132 [ left>> (node>alist) ]
134 [ right>> (node>alist) ]
139 [ root>> (node>alist) ] { } make ;
141 :: (node>subalist-right) ( to-key node end-comparator: ( key1 key2 -- ? ) -- )
143 node key>> to-key end-comparator call :> node-left?
145 node left>> node-left? [ (node>alist) ] [
146 [ to-key ] dip end-comparator (node>subalist-right)
151 right>> [ to-key ] dip
152 end-comparator (node>subalist-right)
155 ] when ; inline recursive
157 :: (node>subalist-left) ( from-key node start-comparator: ( key1 key2 -- ? ) -- )
159 node key>> from-key start-comparator call :> node-right?
163 left>> [ from-key ] dip
164 start-comparator (node>subalist-left)
168 node right>> node-right? [ (node>alist) ] [
169 [ from-key ] dip start-comparator (node>subalist-left)
171 ] when ; inline recursive
173 :: (node>subalist) ( from-key to-key node start-comparator: ( key1 key2 -- ? ) end-comparator: ( key1 key2 -- ? ) -- )
175 node key>> from-key start-comparator call :> node-right?
176 node key>> to-key end-comparator call :> node-left?
179 from-key node left>> node-left?
180 [ start-comparator (node>subalist-left) ]
182 [ to-key ] dip start-comparator
183 end-comparator (node>subalist)
187 node-right? node-left? and [ node entry, ] when
190 to-key node right>> node-right?
191 [ end-comparator (node>subalist-right) ]
193 [ from-key ] 2dip start-comparator
194 end-comparator (node>subalist)
197 ] when ; inline recursive
201 : subtree>alist[) ( from-key to-key tree -- alist )
202 [ root>> [ after=? ] [ before? ] (node>subalist) ] { } make ;
204 : subtree>alist(] ( from-key to-key tree -- alist )
205 [ root>> [ after? ] [ before=? ] (node>subalist) ] { } make ;
207 : subtree>alist[] ( from-key to-key tree -- alist )
208 [ root>> [ after=? ] [ before=? ] (node>subalist) ] { } make ;
210 : subtree>alist() ( from-key to-key tree -- alist )
211 [ root>> [ after? ] [ before? ] (node>subalist) ] { } make ;
213 : headtree>alist[) ( to-key tree -- alist )
214 [ root>> [ before? ] (node>subalist-right) ] { } make ;
216 : headtree>alist[] ( to-key tree -- alist )
217 [ root>> [ before=? ] (node>subalist-right) ] { } make ;
219 : tailtree>alist[] ( from-key tree -- alist )
220 [ root>> [ after=? ] (node>subalist-left) ] { } make ;
222 : tailtree>alist(] ( from-key tree -- alist )
223 [ root>> [ after? ] (node>subalist-left) ] { } make ;
227 : (nodepath-at) ( key node -- )
233 choose-branch (nodepath-at)
237 : nodepath-at ( key tree -- path )
238 [ root>> (nodepath-at) ] { } make ;
240 : right-extremity ( node -- node' )
241 [ dup right>> ] [ nip ] while* ;
243 : left-extremity ( node -- node' )
244 [ dup left>> ] [ nip ] while* ;
246 : lower-node-in-child? ( key node -- ? )
247 [ nip left>> ] [ key>> = ] 2bi and ;
249 : higher-node-in-child? ( key node -- ? )
250 [ nip right>> ] [ key>> = ] 2bi and ;
252 : lower-node ( key tree -- node )
255 reverse 2dup first lower-node-in-child?
256 [ nip first left>> right-extremity ]
257 [ [ key>> after? ] with find nip ] if
260 : higher-node ( key tree -- node )
263 reverse 2dup first higher-node-in-child?
264 [ nip first right>> left-extremity ]
265 [ [ key>> before? ] with find nip ] if
268 : floor-node ( key tree -- node )
269 dupd nodepath-at [ drop f ] [
270 reverse [ key>> after=? ] with find nip
273 : ceiling-node ( key tree -- node )
274 dupd nodepath-at [ drop f ] [
275 reverse [ key>> before=? ] with find nip
278 : first-node ( tree -- node ) root>> dup [ left-extremity ] when ;
280 : last-node ( tree -- node ) root>> dup [ right-extremity ] when ;
284 : lower-entry ( key tree -- pair/f ) lower-node dup [ node>entry ] when ;
286 : higher-entry ( key tree -- pair/f ) higher-node dup [ node>entry ] when ;
288 : floor-entry ( key tree -- pair/f ) floor-node dup [ node>entry ] when ;
290 : ceiling-entry ( key tree -- pair/f ) ceiling-node dup [ node>entry ] when ;
292 : first-entry ( tree -- pair/f ) first-node dup [ node>entry ] when ;
294 : last-entry ( tree -- pair/f ) last-node dup [ node>entry ] when ;
296 : lower-key ( key tree -- key/f ) lower-node dup [ key>> ] when ;
298 : higher-key ( key tree -- key/f ) higher-node dup [ key>> ] when ;
300 : floor-key ( key tree -- key/f ) floor-node dup [ key>> ] when ;
302 : ceiling-key ( key tree -- key/f ) ceiling-node dup [ key>> ] when ;
304 : first-key ( tree -- key/f ) first-node dup [ key>> ] when ;
306 : last-key ( tree -- key/f ) last-node dup [ key>> ] when ;
314 : copy-node-contents ( new old -- new )
316 [ value>> >>value ] bi ;
321 : (prune-extremity) ( parent node -- new-extremity )
323 nipd (prune-extremity)
325 [ delete-node swap set-node-link ] keep
328 : prune-extremity ( node -- new-extremity )
329 ! remove and return the leftmost or rightmost child of this node.
330 ! assumes at least one child
331 dup node-link (prune-extremity) ;
333 : replace-with-child ( node -- node )
334 dup node-link copy-node-contents dup node-link delete-node over set-node-link ;
336 : replace-with-extremity ( node -- node )
337 dup node-link dup node+link [
338 ! predecessor/successor is not the immediate child
339 [ prune-extremity ] with-other-side copy-node-contents
341 ! node-link is the predecessor/successor
342 drop replace-with-child
345 : delete-node-with-two-children ( node -- node )
346 ! randomised to minimise tree unbalancing
347 random-side [ replace-with-extremity ] with-side ;
349 : delete-node ( node -- node )
350 ! delete this node, returning its replacement
351 dup [ right>> ] [ left>> ] bi [
353 drop delete-node-with-two-children
355 nip ! left but no right
358 nip ! right but no left, or no children
361 : delete-bst-node ( key node -- node deleted? )
362 2dup key>> key-side dup 0 eq? [
363 drop nip delete-node t
366 [ node-link delete-bst-node ]
367 [ swap [ set-node-link ] dip ]
373 [ delete-bst-node swap ] change-root
374 swap [ dup dec-count ] when drop ;
379 : clone-nodes ( node -- node' )
382 [ clone-nodes ] change-left
383 [ clone-nodes ] change-right
386 M: tree clone (clone) [ clone-nodes ] change-root ;
388 : ?push-children ( node queue -- )
389 [ [ left>> ] [ right>> ] bi ]
390 [ [ over [ push-front ] [ 2drop ] if ] curry bi@ ] bi* ;
392 : each-bfs-node ( tree quot: ( ... entry -- ... ) -- ... )
393 [ root>> <dlist> [ push-front ] keep dup ] dip
395 [ drop node>entry ] prepose
396 [ ?push-children ] 2bi
397 ] 2curry slurp-deque ; inline
399 : >bfs-alist ( tree -- alist )
400 dup assoc-size <vector> [
401 [ push ] curry each-bfs-node
404 M: tree assoc-clone-like
405 [ dup tree? [ >bfs-alist ] when ] dip call-next-method ;
409 : >tree ( assoc -- tree )
410 T{ tree f f 0 } assoc-clone-like ;
413 \ } [ >tree ] parse-literal ;
417 M: tree assoc-like drop dup tree? [ >tree ] unless ;
419 M: tree assoc-size count>> ;
420 M: tree pprint-delims drop \ TREE{ \ } ;
421 M: tree >pprint-sequence >alist ;
422 M: tree pprint-narrow? drop t ;
424 : node-height ( node -- n )
426 [ left>> ] [ right>> ] bi
427 [ node-height ] bi@ max 1 +
432 : height ( tree -- n )
437 : pop-tree-extremity ( tree node/f -- node/f )
439 [ key>> swap delete-at ] keep node>entry
442 :: slurp-tree ( tree quot: ( ... entry -- ... ) getter: ( tree -- node ) -- ... )
444 [ tree getter call quot call ] until ; inline
448 : pop-tree-left ( tree -- node/f )
449 dup first-node pop-tree-extremity ;
451 : pop-tree-right ( tree -- node/f )
452 dup last-node pop-tree-extremity ;
454 : slurp-tree-left ( tree quot: ( ... entry -- ... ) -- ... )
455 [ pop-tree-left ] slurp-tree ; inline
457 : slurp-tree-right ( tree quot: ( ... entry -- ... ) -- ... )
458 [ pop-tree-right ] slurp-tree ; inline