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 node+link
  27     dup node-link
  28     pick set-node+link
  29     [ set-node-link ] keep ;
  30
  31 : single-rotate ( node -- node )
  32     0 >>balance
  33     0 over node+link
  34     balance<< rotate ;
  35
  36 : pick-balances ( a node -- balance balance )
  37     balance>> {
  38         { [ dup zero? ] [ 2drop 0 0 ] }
  39         { [ 2dup = ] [ nip neg 0 ] }
  40         [ drop 0 swap ]
  41     } cond ;
  42
  43 : double-rotate ( node -- node )
  44     [
  45         node+link [
  46             node-link current-side get neg
  47             over pick-balances rot 0 swap balance<<
  48         ] keep balance<<
  49     ] keep swap >>balance
  50     dup node+link [ rotate ] with-other-side
  51     over set-node+link rotate ;
  52
  53 : select-rotate ( node -- node )
  54     dup node+link balance>> current-side get =
  55     [ double-rotate ] [ single-rotate ] if ;
  56
  57 : balance-insert ( node -- node taller? )
  58     dup balance>> {
  59         { [ dup zero? ] [ drop f ] }
  60         { [ dup abs 2 = ] [ sgn neg [ select-rotate ] with-side f ] }
  61         [ drop t ] ! balance is -1 or 1, tree is taller
  62     } cond ;
  63
  64 DEFER: avl-set
  65
  66 : avl-insert ( value key node -- node taller? created? )
  67     2dup key>> before? left right ? [
  68         [ node-link avl-set ] keep -rot
  69         [ [ set-node-link ] keep ] 2dip swap
  70         [ [ current-side get increase-balance balance-insert ] dip ]
  71         [ f swap ] if
  72     ] with-side ;
  73
  74 : (avl-set) ( value key node -- node taller? created? )
  75     2dup key>> = [
  76         -rot pick key<< >>value f f
  77     ] [ avl-insert ] if ;
  78
  79 : avl-set ( value key node -- node taller? created? )
  80     [ (avl-set) ] [ swap <avl-node> t t ] if* ;
  81
  82 M: avl set-at ( value key node -- )
  83     [ avl-set nip swap ] change-root
  84     swap [ dup inc-count ] when 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 nip swap ] change-root
 152     swap [ dup dec-count ] when drop ;
 153
 154 M: avl new-assoc 2drop <avl> ;
 155
 156 PRIVATE>
 157
 158 : >avl ( assoc -- avl )
 159     T{ avl } assoc-clone-like ;
 160
 161 M: avl assoc-like
 162     drop dup avl? [ >avl ] unless ;
 163
 164 SYNTAX: AVL{
 165     \ } [ >avl ] parse-literal ;
 166
 167 M: avl pprint-delims drop \ AVL{ \ } ;