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 )
30 [ set-node-link ] keep ;
32 : single-rotate ( node -- node )
37 : pick-balances ( a node -- balance balance )
39 { [ dup zero? ] [ 2drop 0 0 ] }
40 { [ 2dup = ] [ nip neg 0 ] }
44 : double-rotate ( node -- node )
47 node-link current-side get neg
48 over pick-balances rot 0 swap balance<<
51 dup node+link [ rotate ] with-other-side
52 over set-node+link rotate ;
54 : select-rotate ( node -- node )
55 dup node+link balance>> current-side get =
56 [ double-rotate ] [ single-rotate ] if ;
58 : balance-insert ( node -- node taller? )
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
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 ]
75 : (avl-set) ( value key node -- node taller? )
77 -rot pick key<< >>value f
80 : avl-set ( value key node -- node taller? )
81 [ (avl-set) ] [ swap <avl-node> t ] if* ;
83 M: avl set-at ( value key node -- )
84 [ avl-set drop ] change-root 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 swap over 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 2drop ] change-root drop ;
153 M: avl new-assoc 2drop <avl> ;
157 : >avl ( assoc -- avl )
158 T{ avl } assoc-clone-like ;
161 drop dup avl? [ >avl ] unless ;
164 \ } [ >avl ] parse-literal ;
166 M: avl pprint-delims drop \ AVL{ \ } ;