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