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
16 TUPLE: avl-node < node balance ;
18 : <avl-node> ( key value -- node )
22 : increase-balance ( node amount -- node )
23 '[ _ + ] change-balance ;
25 : rotate ( node -- node )
29 [ set-node-link ] keep ;
31 : single-rotate ( node -- node )
36 : pick-balances ( a node -- balance balance )
38 { [ dup zero? ] [ 2drop 0 0 ] }
39 { [ 2dup = ] [ nip neg 0 ] }
43 : double-rotate ( node -- node )
46 node-link current-side get neg
47 over pick-balances rot 0 swap balance<<
50 dup node+link [ rotate ] with-other-side
51 over set-node+link rotate ;
53 : select-rotate ( node -- node )
54 dup node+link balance>> current-side get =
55 [ double-rotate ] [ single-rotate ] if ;
57 : balance-insert ( node -- node taller? )
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
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 ]
74 : (avl-set) ( value key node -- node taller? created? )
76 -rot pick key<< >>value f f
79 : avl-set ( value key node -- node taller? created? )
80 [ (avl-set) ] [ swap <avl-node> t t ] if* ;
82 M: avl set-at ( value key node -- )
83 [ avl-set nip swap ] change-root
84 swap [ dup inc-count ] when drop ;
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
94 : rebalance-delete ( node -- node shorter? )
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
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 ]
108 : avl-replace-with-extremity ( to-replace node -- node shorter? )
110 swapd avl-replace-with-extremity [ over set-node-link ] dip
111 [ balance-delete ] [ f ] if
113 [ copy-node-contents drop ] keep node+link t
116 : replace-with-a-child ( node -- node shorter? )
117 ! assumes that node is not a leaf, otherwise will recurse forever
119 dupd [ avl-replace-with-extremity ] with-other-side
120 [ over set-node-link ] dip [ balance-delete ] [ f ] if
122 [ replace-with-a-child ] with-other-side
125 : avl-delete-node ( node -- node shorter? )
126 ! delete this node, returning its replacement, and whether this subtree is
127 ! shorter as a result
131 left [ replace-with-a-child ] with-side
134 GENERIC: avl-delete ( key node -- node shorter? deleted? )
136 M: f avl-delete ( key f -- f f f ) nip f f ;
138 : (avl-delete) ( key node -- node shorter? deleted? )
139 tuck node-link avl-delete [
140 [ over set-node-link ] dip [ balance-delete ] [ f ] if
143 M: avl-node avl-delete ( key node -- node shorter? deleted? )
144 2dup key>> key-side [
145 nip avl-delete-node t
147 [ (avl-delete) ] with-side
150 M: avl delete-at ( key node -- )
151 [ avl-delete nip swap ] change-root
152 swap [ dup dec-count ] when drop ;
154 M: avl new-assoc 2drop <avl> ;
158 : >avl ( assoc -- avl )
159 T{ avl } assoc-clone-like ;
162 drop dup avl? [ >avl ] unless ;
165 \ } [ >avl ] parse-literal ;
167 M: avl pprint-delims drop \ AVL{ \ } ;