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 ;
13 TUPLE: avl-node < node balance ;
15 : <avl-node> ( key value -- node )
19 : increase-balance ( node amount -- )
20 swap [ + ] change-balance drop ;
22 : rotate ( node -- node )
23 dup node+link dup node-link pick set-node+link
26 : single-rotate ( node -- node )
27 0 over (>>balance) 0 over node+link
30 : pick-balances ( a node -- balance balance )
32 { [ dup zero? ] [ 2drop 0 0 ] }
33 { [ over = ] [ neg 0 ] }
37 : double-rotate ( node -- node )
40 node-link current-side get neg
41 over pick-balances rot 0 swap (>>balance)
44 dup node+link [ rotate ] with-other-side
45 over set-node+link rotate ;
47 : select-rotate ( node -- node )
48 dup node+link balance>> current-side get =
49 [ double-rotate ] [ single-rotate ] if ;
51 : balance-insert ( node -- node taller? )
53 { [ dup zero? ] [ drop f ] }
55 [ sgn neg [ select-rotate ] with-side f ] }
56 { [ drop t ] [ t ] } ! balance is -1 or 1, tree is taller
61 : avl-insert ( value key node -- node taller? )
62 2dup key>> before? left right ? [
63 [ node-link avl-set ] keep swap
64 [ tuck set-node-link ] dip
65 [ dup current-side get increase-balance balance-insert ]
69 : (avl-set) ( value key node -- node taller? )
71 -rot pick (>>key) over (>>value) f
74 : avl-set ( value key node -- node taller? )
75 [ (avl-set) ] [ swap <avl-node> t ] if* ;
77 M: avl set-at ( value key node -- node )
78 [ avl-set drop ] change-root drop ;
80 : delete-select-rotate ( node -- node shorter? )
81 dup node+link balance>> zero? [
82 current-side get neg over (>>balance)
83 current-side get over node+link (>>balance) rotate f
88 : rebalance-delete ( node -- node shorter? )
90 { [ dup zero? ] [ drop t ] }
91 { [ dup abs 2 = ] [ sgn neg [ delete-select-rotate ] with-side ] }
92 { [ drop t ] [ f ] } ! balance is -1 or 1, tree is not shorter
95 : balance-delete ( node -- node shorter? )
96 current-side get over balance>> {
97 { [ dup zero? ] [ drop neg over (>>balance) f ] }
98 { [ dupd = ] [ drop 0 >>balance t ] }
99 [ dupd neg increase-balance rebalance-delete ]
102 : avl-replace-with-extremity ( to-replace node -- node shorter? )
104 swapd avl-replace-with-extremity [ over set-node-link ] dip
105 [ balance-delete ] [ f ] if
107 [ copy-node-contents drop ] keep node+link t
110 : replace-with-a-child ( node -- node shorter? )
111 #! assumes that node is not a leaf, otherwise will recurse forever
113 dupd [ avl-replace-with-extremity ] with-other-side
114 [ over set-node-link ] dip [ balance-delete ] [ f ] if
116 [ replace-with-a-child ] with-other-side
119 : avl-delete-node ( node -- node shorter? )
120 #! delete this node, returning its replacement, and whether this subtree is
121 #! shorter as a result
125 left [ replace-with-a-child ] with-side
128 GENERIC: avl-delete ( key node -- node shorter? deleted? )
130 M: f avl-delete ( key f -- f f f ) nip f f ;
132 : (avl-delete) ( key node -- node shorter? deleted? )
133 tuck node-link avl-delete [
134 [ over set-node-link ] dip [ balance-delete ] [ f ] if
137 M: avl-node avl-delete ( key node -- node shorter? deleted? )
138 2dup key>> key-side dup zero? [
139 drop nip avl-delete-node t
141 [ (avl-delete) ] with-side
144 M: avl delete-at ( key node -- )
145 [ avl-delete 2drop ] change-root drop ;
147 M: avl new-assoc 2drop <avl> ;
149 : >avl ( assoc -- avl )
150 T{ avl f f 0 } assoc-clone-like ;
153 drop dup avl? [ >avl ] unless ;
156 \ } [ >avl ] parse-literal ;
158 M: avl pprint-delims drop \ AVL{ \ } ;