aboutsummaryrefslogtreecommitdiffstats
path: root/mcgoron.weight-balanced-trees.internal.scm
blob: 8d9dc6b6210914a6c18715695b2eb3da0e1a34ea (plain) (blame)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
;;; ;;;;;;;;;;;;;;;;;;;
;;; 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)))))

(: join2 (node-type node-type --> node-type))
(define (join2 left right)
  (define split-last
    (the ((struct <wb-tree>) --> node-type *)
         (lambda (tree)
           (with-node (tree data ("<" left) (">" right))
             (if (null? right)
                 (values left data)
                 (let-values (((new-right new-data)
                              (split-last right)))
                   (values (join data left new-right)
                           new-data)))))))
  (if (null? left)
      right
      (let-values (((new-left new-data)
                    (split-last left)))
        (join new-data new-left right))))