aboutsummaryrefslogtreecommitdiff
diff options
context:
space:
mode:
authorMasaya Tojo <masaya@tojo.tokyo>2017-10-26 07:01:45 +0900
committerMasaya Tojo <masaya@tojo.tokyo>2017-10-26 07:01:45 +0900
commit548eaa2f5704967efe0b51ebce26bbece4061574 (patch)
treee88e275ed1d312d9e1248db77f45c5957fa95876
parent0b1cf43d91f2967b9b2576a12b3c79bf0870e242 (diff)
Implement AVL Tree
-rw-r--r--avl.ss312
1 files changed, 312 insertions, 0 deletions
diff --git a/avl.ss b/avl.ss
new file mode 100644
index 0000000..5ae4c85
--- /dev/null
+++ b/avl.ss
@@ -0,0 +1,312 @@
+#!r6rs
+(library (tokyo tojo map avl)
+ (export map? empty? empty
+ insert search (rename (avl:remove remove))
+ size
+ (rename (avl:map map)
+ (avl:filter filter)
+ (avl:partition partition)
+ (avl:for-each for-each)
+ (avl:fold-left fold-left)
+ (avl:fold-right fold-right)))
+ (import (rnrs))
+
+ (define map?
+ (lambda (x)
+ (or (empty? x) (node? x))))
+
+ (define-record-type (avl:empty make-empty empty?)
+ (opaque #t))
+
+ (define error-not-avl-tree
+ (lambda (name x)
+ (unless (map? x)
+ (assertion-violation name "not avl tree" x))))
+
+ (define empty (make-empty))
+
+ (define-record-type (avl:node make-node node?)
+ (fields (immutable kv key&value)
+ (immutable l left)
+ (immutable r right)
+ (immutable h %height))
+ (opaque #t))
+
+ (define height
+ (lambda (tr)
+ (if (empty? tr)
+ 0
+ (%height tr))))
+
+ (define balance-factor
+ (lambda (node)
+ (- (height (left node))
+ (height (right node)))))
+
+ (define height+
+ (lambda (l r)
+ (+ 1 (max (height l)
+ (height r)))))
+
+ ;; (-> (a p (b q c))
+ ;; ((a p b) q c))
+ (define rotate-left
+ (lambda (p)
+ (let ([q (right p)])
+ (let ([a (left p)]
+ [b (left q)]
+ [c (right q)])
+ (let ([p (make-node (key&value p)
+ a b
+ (height+ a b))])
+ (make-node (key&value q)
+ p c
+ (height+ p c)))))))
+
+ ;; (-> ((a p b) q c)
+ ;; (a p (b q c)))
+ (define rotate-right
+ (lambda (q)
+ (let ([p (left q)])
+ (let ([a (left p)]
+ [b (right p)]
+ [c (right q)])
+ (let ([q (make-node (key&value q)
+ b c
+ (height+ b c))])
+ (make-node (key&value p)
+ a q
+ (height+ a q)))))))
+
+
+ ;; (-> ((a p (b q c)) r d)
+ ;; (((a p b) q c) r d) ; left rotation
+ ;; ((a p b) q (c r d))) ; right rotation
+ (define rotate-left-right
+ (lambda (r)
+ (let* ([p (left r)]
+ [q (right p)])
+ (let ([a (left p)]
+ [b (left q)]
+ [c (right q)]
+ [d (right r)])
+ (let ([p (make-node (key&value p)
+ a b
+ (height+ a b))]
+ [r (make-node (key&value r)
+ c d
+ (height+ c d))])
+ (make-node (key&value q)
+ p r
+ (height+ p r)))))))
+
+ ;; (-> (a p ((b q c) r d))
+ ;; (a p (b q (c r d))) ; right rotation
+ ;; ((a p b) q (c r d))) ; left rotation
+ (define rotate-right-left
+ (lambda (p)
+ (let* ([r (right p)]
+ [q (left r)])
+ (let ([a (left p)]
+ [b (left q)]
+ [c (right q)]
+ [d (right r)])
+ (let ([p (make-node (key&value p)
+ a b
+ (height+ a b))]
+ [r (make-node (key&value r)
+ c d
+ (height+ c d))])
+ (make-node (key&value q)
+ p r
+ (height+ p r)))))))
+
+ (define balance
+ (lambda (tr)
+ (if (empty? tr)
+ empty
+ (case (balance-factor tr)
+ [(2)
+ (case (balance-factor (left tr))
+ [(0 1) (rotate-right tr)]
+ [(-1) (rotate-left-right tr)]
+ [else (assertion-violation 'blanace "error")])]
+ [(-2)
+ (case (balance-factor (right tr))
+ [(0 -1) (rotate-left tr)]
+ [(1) (rotate-right-left tr)]
+ [else (assertion-violation 'blanace "error")])]
+ [(-1 0 1) tr]
+ [else
+ (assertion-violation 'balance "error"
+ (balance-factor tr)
+ tr)]))))
+
+ (define insert
+ (case-lambda
+ [(<? =? tr k v)
+ (error-not-avl-tree 'insert tr)
+ (let f ([tr tr])
+ (if (not (or (node? tr) (empty? tr)))
+ (assertion-violation 'insert "not avl tree" tr))
+ (if (empty? tr)
+ (make-node (cons k v) empty empty 0)
+ (let ([kv (key&value tr)])
+ (cond
+ [(=? k (car kv))
+ (make-node (cons k v) (left tr) (right tr) (height tr))]
+ [(<? k (car kv))
+ (let ([l (f (left tr))]
+ [r (right tr)])
+ (balance (make-node (key&value tr)
+ l r
+ (height+ l r))))]
+ [else
+ (let ([l (left tr)]
+ [r (f (right tr))])
+ (balance (make-node (key&value tr)
+ l r
+ (height+ l r))))]))))]
+ [(tr k v)
+ (insert < = tr k v)]))
+
+ (define search
+ (case-lambda
+ [(<? =? tr k)
+ (error-not-avl-tree 'search tr)
+ (let f ([tr tr])
+ (if (empty? tr)
+ #f
+ (let ([kv (key&value tr)])
+ (cond
+ [(=? k (car kv)) (key&value tr)]
+ [(<? k (car kv)) (f (left tr))]
+ [else (f (right tr))]))))]
+ [(tr k) (search < = tr k)]))
+
+ (define node-max
+ (lambda (node)
+ (if (empty? (right node))
+ (values (left node) (key&value node))
+ (let-values ([(tr m) (node-max (right node))])
+ (values (balance
+ (make-node (key&value node)
+ (left node) tr
+ (height+ (left node) tr)))
+ m)))))
+
+ (define rem
+ (lambda (node)
+ (cond
+ [(empty? (left node)) (right node)]
+ [(empty? (right node)) (left node)]
+ [else
+ (let-values ([(tr kv) (node-max (left node))])
+ (balance
+ (make-node kv
+ tr (right node)
+ (height+ tr (right node)))))])))
+
+ (define avl:remove
+ (case-lambda
+ [(<? =? tr k)
+ (error-not-avl-tree 'remove tr)
+ (let f ([tr tr])
+ (if (empty? tr)
+ empty
+ (let ([kv (key&value tr)])
+ (cond
+ [(=? k (car kv)) (rem tr)]
+ [(<? k (car kv))
+ (let ([l (f (left tr))]
+ [r (right tr)])
+ (balance
+ (make-node kv
+ l r
+ (height+ l r))))]
+ [else
+ (let ([l (left tr)]
+ [r (f (right tr))])
+ (balance
+ (make-node kv
+ l r
+ (height+ l r))))]))))]
+ [(tr k)
+ (avl:remove < = tr k)]))
+
+ (define avl:map
+ (lambda (f tr)
+ (error-not-avl-tree 'map tr)
+ (let g ([tr tr])
+ (if (empty? tr)
+ empty
+ (let ([kv (key&value tr)])
+ (make-node (cons (car kv) (f kv))
+ (g (left tr))
+ (g (right tr))
+ (height tr)))))))
+
+ (define avl:for-each
+ (lambda (f tr)
+ (error-not-avl-tree 'for-each tr)
+ (let g [(tr tr)]
+ (if (empty? tr)
+ (if #f #f)
+ (begin
+ (g (left tr))
+ (f (key&value tr))
+ (g (right tr)))))))
+
+ (define fold
+ (lambda (f init tr)
+ (if (empty? tr)
+ init
+ (f (fold f init (left tr))
+ (key&value tr)
+ (fold f init (right tr))))))
+
+ (define avl:fold-left
+ (lambda (f init tr)
+ (error-not-avl-tree 'fold-left tr)
+ (let g [[acc init] (tr tr)]
+ (if (empty? tr)
+ acc
+ (g (f (g acc (left tr))
+ (key&value tr))
+ (right tr))))))
+
+ (define avl:fold-right
+ (lambda (f init tr)
+ (error-not-avl-tree 'fold-right tr)
+ (let g [[acc init] (tr tr)]
+ (if (empty? tr)
+ acc
+ (g (f (key&value tr)
+ (g acc (right tr)))
+ (left tr))))))
+
+ (define size
+ (lambda (tr)
+ (error-not-avl-tree 'size tr)
+ (avl:fold-left (lambda (acc v) (+ acc 1)) 0 tr)))
+
+ (define avl:filter
+ (lambda (p? tr)
+ (error-not-avl-tree 'filter tr)
+ (avl:for-each (lambda (kv)
+ (if (not (p? kv))
+ (set! tr (avl:remove tr (car kv)))))
+ tr)
+ tr))
+
+ (define avl:partition
+ (lambda (p? tr)
+ (error-not-avl-tree 'filter tr)
+ (let ([t tr]
+ [f tr])
+ (avl:for-each (lambda (kv)
+ (if (p? kv)
+ (set! f (avl:remove f (car kv)))
+ (set! t (avl:remove t (car kv)))))
+ tr)
+ (values t f)))))