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 deques dlists kernel make math
   5 math.order namespaces parser prettyprint.custom random sequences
   6 vectors ;
   7 IN: trees
   8
   9 TUPLE: tree root { count integer } ;
  10
  11 <PRIVATE
  12
  13 : new-tree ( class -- tree )
  14     new
  15         f >>root
  16         0 >>count ; inline
  17
  18 PRIVATE>
  19
  20 : <tree> ( -- tree )
  21     tree new-tree ;
  22
  23 INSTANCE: tree assoc
  24
  25 <PRIVATE
  26
  27 TUPLE: node key value left right ;
  28
  29 : new-node ( key value class -- node )
  30     new
  31         swap >>value
  32         swap >>key ; inline
  33
  34 : <node> ( key value -- node )
  35     node new-node ;
  36
  37 SYMBOL: current-side
  38
  39 CONSTANT: left -1
  40 CONSTANT: right 1
  41
  42 : key-side ( k1 k2 -- n )
  43     <=> {
  44         { +lt+ [ left ] }
  45         { +eq+ [ 0 ] }
  46         { +gt+ [ right ] }
  47     } case ;
  48
  49 : go-left? ( -- ? ) current-side get left eq? ;
  50
  51 : inc-count ( tree -- ) [ 1 + ] change-count drop ;
  52
  53 : dec-count ( tree -- ) [ 1 - ] change-count drop ;
  54
  55 : node-link@ ( node ? -- node )
  56     go-left? xor [ left>> ] [ right>> ] if ;
  57
  58 : set-node-link@ ( left parent ? -- )
  59     go-left? xor [ left<< ] [ right<< ] if ;
  60
  61 : node-link ( node -- child ) f node-link@  ;
  62
  63 : set-node-link ( child node -- ) f set-node-link@ ;
  64
  65 : node+link ( node -- child ) t node-link@ ;
  66
  67 : set-node+link ( child node -- ) t set-node-link@ ;
  68
  69 : with-side ( side quot -- )
  70     [ current-side ] dip with-variable ; inline
  71
  72 : with-other-side ( quot -- )
  73     current-side get neg swap with-side ; inline
  74
  75 : go-left ( quot -- ) left swap with-side ; inline
  76
  77 : go-right ( quot -- ) right swap with-side ; inline
  78
  79 : leaf? ( node -- ? )
  80     { [ left>> not ] [ right>> not ] } 1&& ;
  81
  82 : random-side ( -- side )
  83     2 random 0 eq? left right ? ;
  84
  85 : choose-branch ( key node -- key node-left/right )
  86     2dup key>> key-side [ node-link ] with-side ;
  87
  88 : node-at* ( key node -- value ? )
  89     [
  90         2dup key>> = [
  91             nip value>> t
  92         ] [
  93             choose-branch node-at*
  94         ] if
  95     ] [ drop f f ] if* ;
  96
  97 M: tree at*
  98     root>> node-at* ;
  99
 100 : node-set ( value key node -- node new? )
 101     2dup key>> key-side dup 0 eq? [
 102         drop nip swap >>value f
 103     ] [
 104         [
 105             [ node-link [ node-set ] [ swap <node> t ] if* ] keep
 106             swap [ [ set-node-link ] keep ] dip
 107         ] with-side
 108     ] if ;
 109
 110 M: tree set-at
 111     [ [ node-set ] [ swap <node> t ] if* swap ] change-root
 112     swap [ dup inc-count ] when drop ;
 113
 114 : valid-node? ( node -- ? )
 115     [
 116         {
 117             [ dup left>> [ key>> swap key>> before? ] when* ]
 118             [ dup right>> [ key>> swap key>> after? ] when* ]
 119             [ left>> valid-node? ]
 120             [ right>> valid-node? ]
 121         } 1&&
 122     ] [ t ] if* ;
 123
 124 : valid-tree? ( tree -- ? ) root>> valid-node? ;
 125
 126 : node>entry ( node -- entry ) [ key>> ] [ value>> ] bi 2array ;
 127
 128 : entry, ( node -- ) node>entry , ;
 129
 130 : (node>alist) ( node -- )
 131     [
 132         [ left>> (node>alist) ]
 133         [ entry, ]
 134         [ right>> (node>alist) ]
 135         tri
 136     ] when* ;
 137
 138 M: tree >alist
 139     [ root>> (node>alist) ] { } make ;
 140
 141 :: (node>subalist-right) ( to-key node end-comparator: ( key1 key2 -- ? ) -- )
 142     node [
 143         node key>> to-key end-comparator call :> node-left?
 144
 145         node left>> node-left? [ (node>alist) ] [
 146             [ to-key ] dip end-comparator (node>subalist-right)
 147         ] if
 148
 149         node-left? [
 150             node [ entry, ] [
 151                 right>> [ to-key ] dip
 152                 end-comparator (node>subalist-right)
 153             ] bi
 154         ] when
 155     ] when ; inline recursive
 156
 157 :: (node>subalist-left) ( from-key node start-comparator: ( key1 key2 -- ? ) -- )
 158     node [
 159         node key>> from-key start-comparator call :> node-right?
 160
 161         node-right? [
 162             node [
 163                 left>> [ from-key ] dip
 164                 start-comparator (node>subalist-left)
 165             ] [ entry, ] bi
 166         ] when
 167
 168         node right>> node-right? [ (node>alist) ] [
 169             [ from-key ] dip start-comparator (node>subalist-left)
 170         ] if
 171     ] when ; inline recursive
 172
 173 :: (node>subalist) ( from-key to-key node start-comparator: ( key1 key2 -- ? ) end-comparator: ( key1 key2 -- ? ) -- )
 174     node [
 175         node key>> from-key start-comparator call :> node-right?
 176         node key>> to-key end-comparator call :> node-left?
 177
 178         node-right? [
 179             from-key node left>> node-left?
 180             [ start-comparator (node>subalist-left) ]
 181             [
 182                 [ to-key ] dip start-comparator
 183                 end-comparator (node>subalist)
 184             ] if
 185         ] when
 186
 187         node-right? node-left? and [ node entry, ] when
 188
 189         node-left? [
 190             to-key node right>> node-right?
 191             [ end-comparator (node>subalist-right) ]
 192             [
 193                  [ from-key ] 2dip start-comparator
 194                  end-comparator (node>subalist)
 195             ] if
 196         ] when
 197     ] when ; inline recursive
 198
 199 PRIVATE>
 200
 201 : subtree>alist[) ( from-key to-key tree -- alist )
 202     [ root>> [ after=? ] [ before? ] (node>subalist) ] { } make ;
 203
 204 : subtree>alist(] ( from-key to-key tree -- alist )
 205     [ root>> [ after? ] [ before=? ] (node>subalist) ] { } make ;
 206
 207 : subtree>alist[] ( from-key to-key tree -- alist )
 208     [ root>> [ after=? ] [ before=? ] (node>subalist) ] { } make ;
 209
 210 : subtree>alist() ( from-key to-key tree -- alist )
 211     [ root>> [ after? ] [ before? ] (node>subalist) ] { } make ;
 212
 213 : headtree>alist[) ( to-key tree -- alist )
 214     [ root>> [ before? ] (node>subalist-right) ] { } make ;
 215
 216 : headtree>alist[] ( to-key tree -- alist )
 217     [ root>> [ before=? ] (node>subalist-right) ] { } make ;
 218
 219 : tailtree>alist[] ( from-key tree -- alist )
 220     [ root>> [ after=? ] (node>subalist-left) ] { } make ;
 221
 222 : tailtree>alist(] ( from-key tree -- alist )
 223     [ root>> [ after? ] (node>subalist-left) ] { } make ;
 224
 225 <PRIVATE
 226
 227 : (nodepath-at) ( key node -- )
 228     [
 229         dup ,
 230         2dup key>> = [
 231             2drop
 232         ] [
 233             choose-branch (nodepath-at)
 234         ] if
 235     ] [ drop ] if* ;
 236
 237 : nodepath-at ( key tree -- path )
 238     [ root>> (nodepath-at) ] { } make ;
 239
 240 : right-extremity ( node -- node' )
 241     [ dup right>> ] [ nip ] while* ;
 242
 243 : left-extremity ( node -- node' )
 244     [ dup left>> ] [ nip ] while* ;
 245
 246 : lower-node-in-child? ( key node -- ? )
 247     [ nip left>> ] [ key>> = ] 2bi and ;
 248
 249 : higher-node-in-child? ( key node -- ? )
 250     [ nip right>> ] [ key>> = ] 2bi and ;
 251
 252 : lower-node ( key tree -- node )
 253     dupd nodepath-at
 254     [ drop f ] [
 255         reverse 2dup first lower-node-in-child?
 256         [ nip first left>> right-extremity ]
 257         [ [ key>> after? ] with find nip ] if
 258     ] if-empty ;
 259
 260 : higher-node ( key tree -- node )
 261     dupd nodepath-at
 262     [ drop f ] [
 263         reverse 2dup first higher-node-in-child?
 264         [ nip first right>> left-extremity ]
 265         [ [ key>> before? ] with find nip ] if
 266     ] if-empty ;
 267
 268 : floor-node ( key tree -- node )
 269     dupd nodepath-at [ drop f ] [
 270         reverse [ key>> after=? ] with find nip
 271     ] if-empty ;
 272
 273 : ceiling-node ( key tree -- node )
 274     dupd nodepath-at [ drop f ] [
 275         reverse [ key>> before=? ] with find nip
 276     ] if-empty ;
 277
 278 : first-node ( tree -- node ) root>> dup [ left-extremity ] when ;
 279
 280 : last-node ( tree -- node ) root>> dup [ right-extremity ] when ;
 281
 282 PRIVATE>
 283
 284 : lower-entry ( key tree -- pair/f ) lower-node dup [ node>entry ] when ;
 285
 286 : higher-entry ( key tree -- pair/f ) higher-node dup [ node>entry ] when ;
 287
 288 : floor-entry ( key tree -- pair/f ) floor-node dup [ node>entry ] when ;
 289
 290 : ceiling-entry ( key tree -- pair/f ) ceiling-node dup [ node>entry ] when ;
 291
 292 : first-entry ( tree -- pair/f ) first-node dup [ node>entry ] when ;
 293
 294 : last-entry ( tree -- pair/f ) last-node dup [ node>entry ] when ;
 295
 296 : lower-key ( key tree -- key/f ) lower-node dup [ key>> ] when ;
 297
 298 : higher-key ( key tree -- key/f ) higher-node dup [ key>> ] when ;
 299
 300 : floor-key ( key tree -- key/f ) floor-node dup [ key>> ] when ;
 301
 302 : ceiling-key ( key tree -- key/f ) ceiling-node dup [ key>> ] when ;
 303
 304 : first-key ( tree -- key/f ) first-node dup [ key>> ] when ;
 305
 306 : last-key ( tree -- key/f ) last-node dup [ key>> ] when ;
 307
 308 <PRIVATE
 309
 310 M: tree clear-assoc
 311     0 >>count
 312     f >>root drop ;
 313
 314 : copy-node-contents ( new old -- new )
 315     [ key>> >>key ]
 316     [ value>> >>value ] bi ;
 317
 318 ! Deletion
 319 DEFER: delete-node
 320
 321 : (prune-extremity) ( parent node -- new-extremity )
 322     dup node-link [
 323         nipd (prune-extremity)
 324     ] [
 325         [ delete-node swap set-node-link ] keep
 326     ] if* ;
 327
 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) ;
 332
 333 : replace-with-child ( node -- node )
 334     dup node-link copy-node-contents dup node-link delete-node over set-node-link ;
 335
 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
 340     ] [
 341         ! node-link is the predecessor/successor
 342         drop replace-with-child
 343     ] if ;
 344
 345 : delete-node-with-two-children ( node -- node )
 346     ! randomised to minimise tree unbalancing
 347     random-side [ replace-with-extremity ] with-side ;
 348
 349 : delete-node ( node -- node )
 350     ! delete this node, returning its replacement
 351     dup [ right>> ] [ left>> ] bi [
 352         swap [
 353             drop delete-node-with-two-children
 354         ] [
 355             nip ! left but no right
 356         ] if
 357     ] [
 358         nip ! right but no left, or no children
 359     ] if* ;
 360
 361 : delete-bst-node ( key node -- node deleted? )
 362     2dup key>> key-side dup 0 eq? [
 363         drop nip delete-node t
 364     ] [
 365         [
 366             [ node-link delete-bst-node ]
 367             [ swap [ set-node-link ] dip ]
 368             [ swap ] tri
 369         ] with-side
 370     ] if ;
 371
 372 M: tree delete-at
 373     [ delete-bst-node swap ] change-root
 374     swap [ dup dec-count ] when drop ;
 375
 376 M: tree new-assoc
 377     2drop <tree> ;
 378
 379 : clone-nodes ( node -- node' )
 380     dup [
 381         clone
 382         [ clone-nodes ] change-left
 383         [ clone-nodes ] change-right
 384     ] when ;
 385
 386 M: tree clone (clone) [ clone-nodes ] change-root ;
 387
 388 : ?push-children ( node queue -- )
 389     [ [ left>> ] [ right>> ] bi ]
 390     [ [ over [ push-front ] [ 2drop ] if ] curry bi@ ] bi* ;
 391
 392 : each-bfs-node ( tree quot: ( ... entry -- ... ) -- ... )
 393     [ root>> <dlist> [ push-front ] keep dup ] dip
 394     [
 395         [ drop node>entry ] prepose
 396         [ ?push-children ] 2bi
 397     ] 2curry slurp-deque ; inline
 398
 399 : >bfs-alist ( tree -- alist )
 400     dup assoc-size <vector> [
 401         [ push ] curry each-bfs-node
 402     ] keep ;
 403
 404 M: tree assoc-clone-like
 405     [ dup tree? [ >bfs-alist ] when ] dip call-next-method ;
 406
 407 PRIVATE>
 408
 409 : >tree ( assoc -- tree )
 410     T{ tree f f 0 } assoc-clone-like ;
 411
 412 SYNTAX: TREE{
 413     \ } [ >tree ] parse-literal ;
 414
 415 <PRIVATE
 416
 417 M: tree assoc-like drop dup tree? [ >tree ] unless ;
 418
 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 ;
 423
 424 : node-height ( node -- n )
 425     [
 426         [ left>> ] [ right>> ] bi
 427         [ node-height ] bi@ max 1 +
 428     ] [ 0 ] if* ;
 429
 430 PRIVATE>
 431
 432 : height ( tree -- n )
 433     root>> node-height ;
 434
 435 <PRIVATE
 436
 437 : pop-tree-extremity ( tree node/f -- node/f )
 438     dup [
 439         [ key>> swap delete-at ] keep node>entry
 440     ] [ nip ] if ;
 441
 442 :: slurp-tree ( tree quot: ( ... entry -- ... ) getter: ( tree -- node ) -- ... )
 443     [ tree count>> 0 = ]
 444     [ tree getter call quot call ] until ; inline
 445
 446 PRIVATE>
 447
 448 : pop-tree-left ( tree -- node/f )
 449     dup first-node pop-tree-extremity ;
 450
 451 : pop-tree-right ( tree -- node/f )
 452     dup last-node pop-tree-extremity ;
 453
 454 : slurp-tree-left ( tree quot: ( ... entry -- ... ) -- ... )
 455     [ pop-tree-left ] slurp-tree ; inline
 456
 457 : slurp-tree-right ( tree quot: ( ... entry -- ... ) -- ... )
 458     [ pop-tree-right ] slurp-tree ; inline