From 548eaa2f5704967efe0b51ebce26bbece4061574 Mon Sep 17 00:00:00 2001
From: Masaya Tojo <masaya@tojo.tokyo>
Date: Thu, 26 Oct 2017 07:01:45 +0900
Subject: Implement AVL Tree

---
 avl.ss | 312 +++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++
 1 file changed, 312 insertions(+)
 create mode 100644 avl.ss

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)))))
-- 
cgit v1.2.3