application/octet-stream
•
10.53 KB
•
233 lines
;; An avl tree is a balanced binary tree that stores a set of ordered keys.
;; It is represented as an array #[root-node compare-func]
;; An avl node with a key is represented as an array #[height key left-child right-child]
;; An empty avl node is represented as just the symbol :e.
;; The comparison function is called with two keys [cmp a b] and is expected to return
;; -1 if a < b
;; +1 if a > b
;; 0 if a = b.
[def avl/empty :e]
[defn avl/empty? [n] [= :e n]]
[export default-cmp [defn avl/default-cmp [x y] [if [< x y] -1 [if [> x y] 1 0]]]]
[defn avl/typecheck [r k]
[or [avl/empty? [avl/root r]]
[= [type-of k] [type-of [avl/key [avl/root r]]]]
[throw [list :type-error "AVL trees can only contains keys of a single type" k [current-lambda]]]]]
[export tree [defn avl/tree [cmp]
#[avl/empty [or cmp avl/default-cmp]]]]
[defn avl/height [n] [if [avl/empty? n] 0 [array/ref n 0]]]
[defn avl/key [n] [array/ref n 1]]
[defn avl/left [n] [array/ref n 2]]
[defn avl/right [n] [array/ref n 3]]
[defn avl/root [r] [array/ref r 0]]
[defn avl/cmp [r] [array/ref r 1]]
[defn avl/min-node [n]
[if [avl/empty? n]
avl/empty
[let [[l [avl/left n]]]
[if [avl/empty? l]
n
[avl/min-mode l]]]]]
[defn avl/update-left [n l] [array/set! [array/dup n] 2 l]]
[defn avl/update-right [n r] [array/set! [array/dup n] 3 r]]
[defn avl/update-key [n k] [array/set! [array/dup n] 1 k]]
[defn avl/update-root [t r] [array/set! [array/dup t] 0 r]]
[defn avl/update-height [n]
[array/set! [array/dup n] 0
[+ 1 [max [avl/height [avl/left n]]
[avl/height [avl/right n]]]]]]
;; y x
;; / \ / \
;; / \ / \
;; x T3 T1 y
;; / \ / \
;; / \ / \
;; T1 T2 T2 T3
;;
;; rotate-left transforms the left node configuration to the right.
;; rotate-right transforms the right node configuration to the left.
[defn avl/rotate-right [y]
[let [[x [avl/left y]]]
[avl/update-height [avl/update-right x [avl/update-height [avl/update-left y [avl/right x]]]]]]]
[defn avl/rotate-left [x]
[let [[y [avl/right x]]]
[avl/update-height [avl/update-left y [avl/update-height [avl/update-right x [avl/left y]]]]]]]
[defn avl/balance [n]
[if [avl/empty? n]
0
[- [avl/height [avl/left n]] [avl/height [avl/right n]]]]]
[defn avl/insert-rebalance [n cmp v]
[let [[b [avl/balance n]]]
[cond
[[> b 1]
[case [cmp v [avl/key [avl/left n]]]
[-1 [avl/rotate-right n]]
[1 [avl/rotate-right [avl/update-left n [avl/rotate-left [avl/left n]]]]]
[0 n]]]
[[< b -1]
[case [cmp v [avl/key [avl/right n]]]
[1 [avl/rotate-left n]]
[-1 [avl/rotate-left [avl/update-right n [avl/rotate-right [avl/right n]]]]]
[0 n]]]
[#t n]]]]
[defn avl/node-insert [n cmp v]
[if [avl/empty? n]
#[1 v avl/empty avl/empty]
[case [cmp v [avl/key n]]
[-1 [avl/insert-rebalance [avl/update-height [avl/update-left n [avl/node-insert [avl/left n] cmp v]]] cmp v]]
[1 [avl/insert-rebalance [avl/update-height [avl/update-right n [avl/node-insert [avl/right n] cmp v]]] cmp v]]
[0 [avl/update-key n v]]]]]
[export insert [defn avl/insert [t v]
"Insert key V into tree T. If a node with an equivalent key already exists, its key is updated to V"
[avl/typecheck t v]
[avl/update-root t [avl/node-insert [avl/root t] [avl/cmp t] v]]]]
[defn avl/node-get [n cmp v]
[if [avl/empty? n]
#nil
[case [cmp v [avl/key n]]
[0 [avl/key n]]
[-1 [avl/node-get [avl/left n] cmp v]]
[1 [avl/node-get [avl/right n] cmp v]]]]]
[export get [defn avl/get [t v]
"Retrieve the key V from tree T, or #nil if V is not in it"
[if [or [avl/empty? [avl/root t]]
[not= [type-of v] [type-of [avl/key [avl/root t]]]]]
#nil
[avl/node-get [avl/root t] [avl/cmp t] v]]]]
[export from-list [defn avl/from-list [l cmp]
"Create a new avl tree using the keys in L and the comparison function CMP"
[list/reduce l avl/insert [avl/tree cmp]]]]
[defn avl/remove-rebalance [n]
[if [avl/empty? n]
n
[let [[b [avl/balance n]]
[l [avl/left n]]
[r [avl/right n]]]
[cond
[[> b 1]
[if [>= [avl/balance l] 0]
[avl/rotate-right n]
[avl/rotate-right [avl/update-left n [avl/rotate-left l]]]]]
[[< b -1]
[if [<= [avl/balance r] 0]
[avl/rotate-left n]
[avl/rotate-left [avl/update-right n [avl/rotate-right r]]]]]
[#t n]]]]]
[defn avl/node-remove [n cmp v]
[if [avl/empty? n]
n
[let [[root
[case [cmp v [avl/key n]]
[-1 [avl/update-left n [avl/node-remove [avl/left n] cmp v]]]
[1 [avl/update-right n [avl/node-remove [avl/right n] cmp v]]]
[0 [cond
[[avl/empty? [avl/left n]] [avl/right n]]
[[avl/empty? [avl/right n]] [avl/left n]]
[#t [let [[k [avl/key [avl/min-node [avl/right n]]]]]
[avl/update-key [avl/update-right [avl/right n] [avl/node-remove [avl/right n] cmp v]] k]]]]]]]]
[set! root [avl/update-height root]]
[avl/remove-rebalance root]]]]
[export remove [defn avl/remove [t v]
"Remove the key V from tree T if it is contained within it"
[avl/update-root t [avl/node-remove [avl/root t] [avl/cmp t] v]]]]
[defn avl/equal-node? [a b]
[if [avl/empty? a]
[avl/empty? b]
[and [equal? [avl/key a] [avl/key b]]
[avl/equal-node? [avl/left a] [avl/left b]]
[avl/equal-node? [avl/right a] [avl/right b]]]]]
[export equal? [defn avl/equal? [a b]
"Test if two avl trees are equal"
[avl/equal-node? [avl/root a] [avl/root b]]]]
[defn avl/reduce-node [node o s]
[if [avl/empty? node]
s
[o [avl/key node]
[avl/reduce-node [avl/right node] o
[avl/reduce-node [avl/left node] o s]]]]]
[export reduce [defn avl/reduce [t o s]
"Reduce T in-order with a reducer O taking a key and the result of the reductions of one subtree"
[avl/reduce-node [avl/root t] o s]]]
[defn avl/reduce-node-bin [n o s]
[if [avl/empty? n]
s
[o [o [avl/key n]
[avl/reduce-node-bin [avl/left n] o s]]
[avl/reduce-node-bin [avl/right n] o s]]]]
[export reduce-bin [defn avl/reduce-bin [t o s]
"Reduce T with a reducer O taking a key and the result of the reductions of both subtrees"
[avl/reduce-node-bin [avl/root t] o s]]]
[export map [defn avl/map [t f]
"Create a new avl tree by mapping each key in T using F, using the same comparison function as T"
[avl/reduce t [fn [x acc] [avl/insert acc [f x]]] [avl/tree [avl/cmp t]]]]]
[export map-to [defn avl/map-to [t f cmp]
"Create a new avl tree by mapping each key in in T using F, using the comparison function CMP, which may be different from the comparison used in T"
[avl/reduce t [fn [x acc] [avl/insert acc [f x]]] [avl/tree cmp]]]]
[export to-list [defn avl/to-list [t]
[avl/reduce t cons #nil]]]
[deftest 1 [avl/get [avl/from-list '[1 2 3]] 1]]
[deftest 4 [avl/get [avl/from-list '[1 2 3 4 5]] 4]]
[deftest 15 [avl/get [avl/from-list '[23 42 100 10 15 64 101]] 15]]
[deftest 100 [avl/get [avl/from-list '[23 42 100 100 101 100]] 100]]
[deftest 101 [avl/get [avl/from-list '[23 42 100 100 101 100]] 101]]
[deftest 100 [avl/get [avl/remove [avl/from-list '[23 42 100 10 64 101]] 15] 100]]
[deftest #nil [avl/get [avl/from-list '[1 2 3]] 4]]
[deftest #nil [avl/get [avl/from-list #nil] 1]]
[deftest #nil [avl/get [avl/remove [avl/from-list '[23 42 100 10 15 64 101]] 15] 15]]
[deftest #nil [avl/get [avl/remove [avl/from-list '[23 42 100 10 64 101]] 15] 15]]
[deftest 355 [avl/reduce [avl/from-list '[23 42 100 10 15 64 101]] + 0]]
[deftest 355 [avl/reduce-bin [avl/from-list '[23 42 100 10 15 64 101]] + 0]]
[deftest -19 [avl/reduce [avl/from-list '[23 42 100 10 15 64 101]] - 0]]
[deftest 125 [avl/reduce-bin [avl/from-list '[23 42 100 10 15 64 101]] - 0]]
[deftest '[42 100 101 64 15 23 10] [avl/to-list [avl/from-list '[23 42 100 10 15 64 101]]]]
[deftest 200 [avl/get [avl/map [avl/from-list '[23 42 100 10 15 64 101]] [fn [x] [* x 2]]] 200]]
[deftest #nil [avl/get [avl/map [avl/from-list '[23 42 100 10 15 64 101]] [fn [x] [* x 2]]] 100]]
[deftest 100 [avl/get [avl/insert [avl/map [avl/remove [avl/from-list '[23 42 100 10 15 64 101]] 100] [fn [x] [* x 2]]] 100] 100]]
[deftest "100" [avl/get [avl/map [avl/from-list '[23 42 100 10 15 64 101]] string] "100"]]
[deftest "100" [avl/get [avl/map-to [avl/from-list '[23 42 100 10 15 64 101]] string [fn [x y] [if [< x y] -1 [if [> x y] 1 0]]]] "100"]]
[deftest "a" [avl/get [avl/from-list '["a" "b" "c"]] "a"]]
[deftest "ein" [avl/get [avl/from-list '["Dies" "ist" "ein" "Test"]] "ein"]]
[deftest #nil [avl/get [avl/from-list '["a" "b" "c"]] "d"]]
[deftest #nil [avl/get [avl/from-list '["a" "b" "c"]] :a]]
[deftest :c [avl/get [avl/from-list '[:a :b :c]] :c]]
[deftest #nil [avl/get [avl/from-list '[:a :b :c]] :d]]
[deftest #nil [avl/get [avl/from-list '[:a :b :c]] "c"]]
[deftest 'c [avl/get [avl/from-list '[c b a]] 'c]]
[deftest 'c [avl/get [avl/from-list '[b c a]] 'c]]
[deftest 'c [avl/get [avl/from-list '[b c b b b a]] 'c]]
[deftest '[b c a] [avl/to-list [avl/from-list '[b c b b b a]] 'c]]
[deftest #nil [avl/get [avl/from-list '[c b a]] 'd]]
[deftest :type-error [try car [avl/from-list '[c :b a]]]]
[deftest :type-error [try car [avl/from-list '["c" b a]]]]
[deftest :type-error [try car [avl/from-list '[123 b a]]]]
