succesfully test join

1 files changed, 221 insertions, 0 deletions

diff --git a/mcgoron.weight-balanced-trees.scm b/mcgoron.weight-balanced-trees.scm
new file mode 100644
index 0000000..a3c9059
--- /dev/null
+++ b/mcgoron.weight-balanced-trees.scm
@@ -0,0 +1,221 @@
+;;; ;;;;;;;;;;;;;;;;;;;
+;;; Definition of nodes and functions to calculate values for nodes.
+;;; ;;;;;;;;;;;;;;;;;;;
+
+(define-type node-type (or (struct <wb-tree>) null))
+
+(: %wb-tree-node (* fixnum node-type node-type --> (struct <wb-tree>)))
+(: get-data ((struct <wb-tree>) --> *))
+(: %get-weight ((struct <wb-tree>) --> fixnum))
+(: get-left ((struct <wb-tree>) --> node-type))
+(: get-right ((struct <wb-tree>) --> node-type))
+
+(define-record-type <wb-tree>
+  (%wb-tree-node data weight left right)
+  non-null-wb-tree-node?
+  (data get-data)
+  (weight %get-weight)
+  (left get-left)
+  (right get-right))
+
+(define (wb-tree-node? x)
+  (or (null? x) (non-null-wb-tree-node? x)))
+
+(: get-weight (node-type --> fixnum))
+(define (get-weight node)
+  ;; Get the stored size of a node.
+  (cond
+    ((null? node) 1)
+    (else (%get-weight node))))
+
+(: get-size (node-type --> fixnum))
+(define (get-size node)
+  (- (get-weight node) 1))
+
+(: fixnum-calculate-weight (fixnum fixnum --> fixnum))
+(define (fixnum-calculate-weight left-weight right-weight)
+  ;; Calculate the weight of a node given the weight of its children.
+  (+ left-weight right-weight))
+
+(: calculate-weight (node-type node-type --> fixnum))
+(define (calculate-weight left right)
+  ;; Calculate the weight of a node that has children `left` and `right`.
+  (fixnum-calculate-weight (get-weight left) (get-weight right)))
+
+(: wb-tree-node (* node-type node-type --> (struct <wb-tree>)))
+(define (wb-tree-node data left right)
+  ;; Construct a node with `data`, `left`, and `right`, with the correct
+  ;; weight.
+  (%wb-tree-node data (calculate-weight left right) left right))
+
+(: balanced-as-child? (fixnum fixnum --> boolean))
+(define (balanced-as-child? child-weight node-weight)
+  ;; Determine if child would be weight-balanced if its parent was a node
+  ;; with weight `node-weight`.
+  (let ((alpha #e0.29))
+    (>= child-weight (* alpha node-weight))))
+
+(: fixnum-would-be-balanced? (fixnum fixnum --> boolean))
+(define (fixnum-would-be-balanced? left-weight right-weight)
+  ;; Determine if the two weights would be balanced if placed into a node.
+  (let ((size (+ left-weight right-weight)))
+    (and (balanced-as-child? left-weight size)
+         (balanced-as-child? right-weight size))))
+
+(: would-be-balanced? (node-type node-type --> boolean))
+(define (would-be-balanced? left right)
+  ;; Determine if the two nodes would be balanced if placed into a node.
+  (fixnum-would-be-balanced? (get-weight left) (get-weight right))) 
+
+(: heavy<=> (node-type node-type --> fixnum))
+(define (heavy<=> left right)
+  ;; Return 1 if right > left, -1 if left < right, and 0 if left = right
+  ;; weightwise.
+  (let ((left (get-weight left))
+        (right (get-weight right)))
+    (cond
+      ((< left right) -1)
+      ((> left right) 1)
+      (else 0))))
+
+(: balanced? (node-type --> boolean))
+(define (balanced? node)
+  ;; Recursively check if node is weight balanced.
+  (cond
+    ((null? node) #t)
+    (else (let ((left (get-left node))
+                (right (get-right node))
+                (weight (get-weight node)))
+            (and (balanced? left) (balanced? right)
+                 (balanced-as-child? (get-weight left) weight)
+                 (balanced-as-child? (get-weight right) weight))))))
+
+;;; ;;;;;;;;;;;;;;;;;;;;;;;;;
+;;; Convert in-order vectors to ordered trees.
+;;; ;;;;;;;;;;;;;;;;;;;;;;;;;
+
+(: in-order-vector->node (vector --> node-type))
+(define (in-order-vector->node vec)
+  (define (divide left right)
+    (if (< right left)
+        '()
+        (let ((midpoint (floor (/ (+ left right) 2))))
+          (wb-tree-node (vector-ref vec midpoint)
+                        (divide left (- midpoint 1))
+                        (divide (+ midpoint 1) right)))))
+  (divide 0 (- (vector-length vec) 1)))
+
+(: node->in-order-list (node-type --> list))
+(define (node->in-order-list node)
+  (if (null? node)
+      '()
+      (append (node->in-order-list (get-left node))
+              (list (get-data node))
+              (node->in-order-list (get-right node)))))
+
+;;; ;;;;;;;;;;;;;;;;;;;;
+;;; Fundamental tree operations
+;;; 
+;;; These macros are used to automatically generate symmetric cases for
+;;; tree operations. They work with a trick with the syntax-rules
+;;; matcher where it will match literal strings.
+;;; 
+;;; Since syntax-rules is not eager it cannot calculate the inverse
+;;; direction, so both must be supplied. This is fine for most cases,
+;;; and helps make the algorithm more explicit.
+;;; ;;;;;;;;;;;;;;;;;;;;
+
+(define-syntax s-make-node
+  ;; Create a node with these directions.
+  (syntax-rules ()
+    ((_ data ("<" left) (">" right))
+     (wb-tree-node data left right))
+    ((_ data (">" right) ("<" left))
+     (s-make-node data ("<" left) (">" right)))))
+
+(define-syntax with-node
+  (syntax-rules ()
+    ((with-node (%node data ("<" left) (">" right)) body ...)
+     (let* ((node %node)
+            (left (get-left node))
+            (right (get-right node))
+            (data (get-data node)))
+       body ...))
+    ((with-node (%node data (">" right) ("<" left)) body ...)
+     (with-node (%node data ("<" left) (">" right)) body ...))))
+
+(define-syntax s-get
+  (syntax-rules ()
+    ((s-get "<" node) (get-left node))
+    ((s-get ">" node) (get-right node))))
+
+(define-syntax s-rotate
+  ;; Generate rotation based on direction. Rotations are:
+  ;; 
+  ;;     A                  C
+  ;;    / \                / \
+  ;;   B   C      ->      A   E
+  ;;      / \            / \
+  ;;     D   E          B   D
+  ;; 
+  ;;     A                  C
+  ;;    / \                / \
+  ;;   C   B      ->      E   A
+  ;;  / \                    / \
+  ;; E   D                  D   B
+  (syntax-rules ()
+    ((_ dir invdir A)
+     (with-node (A A-data (dir B) (invdir C))
+       (with-node (C C-data (dir D) (invdir E))
+         (s-make-node C-data
+                      (dir (s-make-node A-data (dir B) (invdir D)))
+                      (invdir E)))))))
+
+(define-syntax s-join
+  ;; Generate a macro that traverses `invdir` to make a balanced tree with
+  ;; `dir` in it and `data` in the middle.
+  (syntax-rules ()
+    ((s-join %data (dir init-in-dir) (invdir init-in-invdir))
+     (let ((in-dir init-in-dir)
+           (%data data))
+       (let join ((in-invdir init-in-invdir))
+         (if (would-be-balanced? in-invdir in-dir)
+             (s-make-node data
+                          (invdir in-invdir)
+                          (dir in-dir))
+             (with-node (in-invdir invdir-data
+                                   (dir dir-in-invdir)
+                                   (invdir invdir-in-invdir))
+               (let ((new-dir (join dir-in-invdir)))
+                 (if (would-be-balanced? invdir-in-invdir new-dir)
+                     (s-make-node invdir-data
+                                  (invdir invdir-in-invdir)
+                                  (dir new-dir))
+                     (with-node (new-dir _
+                                         (dir dir-in-new-dir)
+                                         (invdir invdir-in-new-dir))
+                       (if (and (would-be-balanced? invdir-in-invdir
+                                                    invdir-in-new-dir)
+                                (fixnum-would-be-balanced?
+                                 (+ (get-weight invdir-in-invdir)
+                                    (get-weight invdir-in-new-dir))
+                                 (get-weight dir-in-new-dir)))
+                           (s-rotate invdir
+                                     dir
+                                     (s-make-node invdir-data
+                                                  (invdir invdir-in-invdir)
+                                                  (dir new-dir)))
+                           (s-rotate invdir
+                                     dir
+                                     (s-make-node invdir-data
+                                                  (invdir invdir-in-invdir)
+                                                  (dir (s-rotate dir invdir new-dir)))))))))))))))
+
+(: join (* node-type node-type --> (struct <wb-tree>)))
+(define (join data left right)
+  (let ((dir (heavy<=> left right)))
+    (cond
+      ((positive? dir) (s-join data (">" right) ("<" left)))
+      ((negative? dir) (s-join data ("<" left) (">" right)))
+      (else (wb-tree-node data left right)))))
+