extra/trees/avl/avl.factor

   1 ! Copyright (C) 2007 Alex Chapman
   2 ! See http://factorcode.org/license.txt for BSD license.
   3 USING: combinators kernel generic math math.functions
   4 math.parser namespaces io sequences trees shuffle
   5 assocs parser accessors math.order prettyprint.custom ;
   6 IN: trees.avl
   7
   8 TUPLE: avl < tree ;
   9
  10 : <avl> ( -- tree )
  11     avl new-tree ;
  12
  13 TUPLE: avl-node < node balance ;
  14
  15 : <avl-node> ( key value -- node )
  16     avl-node new-node
  17         0 >>balance ;
  18
  19 : increase-balance ( node amount -- )
  20     swap [ + ] change-balance drop ;
  21
  22 : rotate ( node -- node )
  23     dup node+link dup node-link pick set-node+link
  24     tuck set-node-link ;
  25
  26 : single-rotate ( node -- node )
  27     0 over (>>balance) 0 over node+link
  28     (>>balance) rotate ;
  29
  30 : pick-balances ( a node -- balance balance )
  31     balance>> {
  32         { [ dup zero? ] [ 2drop 0 0 ] }
  33         { [ over = ] [ neg 0 ] }
  34         [ 0 swap ]
  35     } cond ;
  36
  37 : double-rotate ( node -- node )
  38     [
  39         node+link [
  40             node-link current-side get neg
  41             over pick-balances rot 0 swap (>>balance)
  42         ] keep (>>balance)
  43     ] keep swap >>balance
  44     dup node+link [ rotate ] with-other-side
  45     over set-node+link rotate ;
  46
  47 : select-rotate ( node -- node )
  48     dup node+link balance>> current-side get =
  49     [ double-rotate ] [ single-rotate ] if ;
  50
  51 : balance-insert ( node -- node taller? )
  52     dup balance>> {
  53         { [ dup zero? ] [ drop f ] }
  54         { [ dup abs 2 = ]
  55           [ sgn neg [ select-rotate ] with-side f ] }
  56         { [ drop t ] [ t ] } ! balance is -1 or 1, tree is taller
  57     } cond ;
  58
  59 DEFER: avl-set
  60
  61 : avl-insert ( value key node -- node taller? )
  62     2dup key>> before? left right ? [
  63         [ node-link avl-set ] keep swap
  64         [ tuck set-node-link ] dip
  65         [ dup current-side get increase-balance balance-insert ]
  66         [ f ] if
  67     ] with-side ;
  68
  69 : (avl-set) ( value key node -- node taller? )
  70     2dup key>> = [
  71         -rot pick (>>key) over (>>value) f
  72     ] [ avl-insert ] if ;
  73
  74 : avl-set ( value key node -- node taller? )
  75     [ (avl-set) ] [ swap <avl-node> t ] if* ;
  76
  77 M: avl set-at ( value key node -- node )
  78     [ avl-set drop ] change-root drop ;
  79
  80 : delete-select-rotate ( node -- node shorter? )
  81     dup node+link balance>> zero? [
  82         current-side get neg over (>>balance)
  83         current-side get over node+link (>>balance) rotate f
  84     ] [
  85         select-rotate t
  86     ] if ;
  87
  88 : rebalance-delete ( node -- node shorter? )
  89     dup balance>> {
  90         { [ dup zero? ] [ drop t ] }
  91         { [ dup abs 2 = ] [ sgn neg [ delete-select-rotate ] with-side ] }
  92         { [ drop t ] [ f ] } ! balance is -1 or 1, tree is not shorter
  93     } cond ;
  94
  95 : balance-delete ( node -- node shorter? )
  96     current-side get over balance>> {
  97         { [ dup zero? ] [ drop neg over (>>balance) f ] }
  98         { [ dupd = ] [ drop 0 >>balance t ] }
  99         [ dupd neg increase-balance rebalance-delete ]
 100     } cond ;
 101
 102 : avl-replace-with-extremity ( to-replace node -- node shorter? )
 103     dup node-link [
 104         swapd avl-replace-with-extremity [ over set-node-link ] dip
 105         [ balance-delete ] [ f ] if
 106     ] [
 107         [ copy-node-contents drop ] keep node+link t
 108     ] if* ;
 109
 110 : replace-with-a-child ( node -- node shorter? )
 111     #! assumes that node is not a leaf, otherwise will recurse forever
 112     dup node-link [
 113         dupd [ avl-replace-with-extremity ] with-other-side
 114         [ over set-node-link ] dip [ balance-delete ] [ f ] if
 115     ] [
 116         [ replace-with-a-child ] with-other-side
 117     ] if* ;
 118
 119 : avl-delete-node ( node -- node shorter? )
 120     #! delete this node, returning its replacement, and whether this subtree is
 121     #! shorter as a result
 122     dup leaf? [
 123         drop f t
 124     ] [
 125         left [ replace-with-a-child ] with-side
 126     ] if ;
 127
 128 GENERIC: avl-delete ( key node -- node shorter? deleted? )
 129
 130 M: f avl-delete ( key f -- f f f ) nip f f ;
 131
 132 : (avl-delete) ( key node -- node shorter? deleted? )
 133     tuck node-link avl-delete [
 134         [ over set-node-link ] dip [ balance-delete ] [ f ] if
 135     ] dip ;
 136
 137 M: avl-node avl-delete ( key node -- node shorter? deleted? )
 138     2dup key>> key-side dup zero? [
 139         drop nip avl-delete-node t
 140     ] [
 141         [ (avl-delete) ] with-side
 142     ] if ;
 143
 144 M: avl delete-at ( key node -- )
 145     [ avl-delete 2drop ] change-root drop ;
 146
 147 M: avl new-assoc 2drop <avl> ;
 148
 149 : >avl ( assoc -- avl )
 150     T{ avl f f 0 } assoc-clone-like ;
 151
 152 M: avl assoc-like
 153     drop dup avl? [ >avl ] unless ;
 154
 155 SYNTAX: AVL{
 156     \ } [ >avl ] parse-literal ;
 157
 158 M: avl pprint-delims drop \ AVL{ \ } ;