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 ;
   7 IN: trees.avl
   8
   9 TUPLE: avl < tree ;
  10
  11 : <avl> ( -- tree )
  12     avl new-tree ;
  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 ;
  21
  22 : increase-balance ( node amount -- )
  23     swap [ + ] change-balance drop ;
  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         { [ over = ] [ neg 0 ] }
  40         [ 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 = ]
  61           [ sgn neg [ select-rotate ] with-side f ] }
  62         { [ drop t ] [ 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         [ dup 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<< over 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 -- 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 over 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 t ] [ 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 over balance<< f ] }
 104         { [ dupd = ] [ drop 0 >>balance t ] }
 105         [ dupd 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     tuck 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 dup zero? [
 145         drop nip avl-delete-node t
 146     ] [
 147         [ (avl-delete) ] with-side
 148     ] if ;
 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 f f 0 } 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{ \ } ;