source: project/binary-tree/avltree.scm @ 5442

Last change on this file since 5442 was 5442, checked in by Kon Lovett, 13 years ago

Chg for misc-extn 3.0

File size: 16.9 KB
Line 
1;;;; avltree.scm
2;;;; Kon Lovett, Oct '06
3;;;; Stephen J. Bevan   <bevan@cs.man.ac.uk> Oct 23 1993
4
5;; Issues
6;;
7;; - Only re-uses deleted nodes when keys match
8
9(use srfi-1)
10(use misc-extn-record)
11
12(eval-when (compile)
13        (declare
14                (usual-integrations)
15                (fixnum)
16                (inline)
17    (no-procedure-checks)
18    (no-bound-checks)
19                (export
20                        make-avltree
21                        alist->avltree
22                        avltree->alist
23                        avltree?
24                        avltree-empty?
25                        avltree-less-than-function
26                        avltree-size
27                        avltree-keys
28                        avltree-values
29                        avltree-ref
30                        avltree-ref/default
31                        avltree-exists?
32                        avltree-update!
33                        avltree-update!/default
34                        avltree-set!
35                        avltree-delete!
36                        avltree-vacuum!
37                        avltree-merge!
38                        avltree-update
39                        avltree-update/default
40                        avltree-set
41                        avltree-delete
42                        avltree-copy
43                        avltree-merge
44                        avltree-walk
45                        avltree-fold
46                        avltree-enfold
47                        avltree-bifold) ) )
48
49;;;
50
51(define true (lambda _ #t))
52(define false (lambda _ #f))
53
54(define-inline-unchecked-record-type avlnode
55        (%make-node left key value deleted height right)
56        %node?
57        (deleted %node-deleted? %node-deleted-set!)
58        (left %node-left %node-left-set!)
59        (right %node-right %node-right-set!)
60        (height %node-height %node-height-set!)
61        (key %node-key %node-key-set!)
62        (value %node-value %node-value-set!) )
63
64(define-record-printer (avlnode nd out)
65        (fprintf out "(node ~A ~A ~S ~S)"
66                (%node-deleted? nd)
67                (%node-height nd)
68                (%node-key nd)
69                (%node-value nd)) )
70
71(define *empty-node* '())
72
73(define %empty-node? null?)
74
75(define (make-node left key value height right)
76        (%make-node left key value #f height right) )
77
78(define-inline (make-leaf k v)
79        (make-node *empty-node* k v 1 *empty-node*) )
80
81(define-inline (*node-height nd)
82        (if (%empty-node? nd) 0 (%node-height nd) ) )
83
84;;;
85
86(define-inline (avltree::max a b)
87  (+ 1 (max a b)) )
88
89;; Insert an element with the given KEY into TREE.  If an element with
90;; the same key is already in the tree, IF-FOUND is called with the
91;; element already there and it should return either it or the element
92;; to be inserted.  If an element with the same key is not in TREE,
93;; then IF-NOT-FOUND is called with no arguments and this should
94;; return the element to be inserted.
95;;-doc
96
97(define (avltree::combine left-tree ak av ad middle-tree ck cv cd right-tree)
98        (let ([left-tree-height (*node-height left-tree)]
99                                [middle-tree-height (*node-height middle-tree)]
100                                [right-tree-height (*node-height right-tree)])
101                (cond
102                        [(and (> middle-tree-height left-tree-height)
103                                                (> middle-tree-height right-tree-height))
104                                (%make-node
105                                        (%make-node left-tree ak av ad (+ 1 left-tree-height) (%node-left middle-tree))
106                                        (%node-key middle-tree) (%node-value middle-tree) (%node-deleted? middle-tree)
107                                        (+ 2 left-tree-height)
108                                        (%make-node (%node-right middle-tree) ck cv cd (+ 1 right-tree-height) right-tree))]
109                        [(and (>= left-tree-height middle-tree-height)
110                                                (>= left-tree-height right-tree-height))
111                                (let* ([middle-right-max (avltree::max middle-tree-height right-tree-height)]
112                                                         [left-middle-right-max (avltree::max middle-right-max left-tree-height)])
113                                        (%make-node
114                                                left-tree
115                                                ak av ad
116                                                left-middle-right-max
117                                                (%make-node middle-tree ck cv cd middle-right-max right-tree)))]
118                        [else
119                                (let* ([left-middle-max (avltree::max left-tree-height middle-tree-height)]
120                                                         [left-middle-right-max (avltree::max left-middle-max right-tree-height)])
121                                        (%make-node
122                                                (%make-node left-tree ak av ad left-middle-max middle-tree)
123                                                ck cv cd
124                                                left-middle-right-max
125                                                right-tree))]) ) )
126
127(define (avltree::add t k less-than if-found if-not-found)
128        (let loop ([t t])
129                (if (%empty-node? t)
130                        (if-not-found)
131                        (let ([ck (%node-key t)])
132                                (cond
133                                        [(less-than ck k)
134                                                (let ([n (loop (%node-right t))])
135                                                        (avltree::combine
136                                                                (%node-left t)
137                                                                (%node-key t) (%node-value t) (%node-deleted? t)
138                                                                (%node-left n)
139                                                                (%node-key n) (%node-value n) (%node-deleted? n)
140                                                                (%node-right n)))]
141                                        [(less-than k ck)
142                                                (let ([n (loop (%node-left t))])
143                                                        (avltree::combine
144                                                                (%node-left n)
145                                                                (%node-key n) (%node-value n) (%node-deleted? n)
146                                                                (%node-right n)
147                                                                (%node-key t) (%node-value t) (%node-deleted? t)
148                                                                (%node-right t)))]
149                                        [else
150                                                (if-found t)] ) ) ) ) )
151
152(define (avltree:merge-insert t k less-than if-found if-not-found)
153        (let ([merge
154                                        (lambda (t)
155                                                (make-node
156                                                        (%node-left t)
157                                                        k
158                                                        (if (%node-deleted? t) (if-not-found) (if-found (%node-value t)))
159                                                        (%node-height t)
160                                                        (%node-right t)))]
161                                [add
162                                        (lambda ()
163                                                (make-leaf k (if-not-found)))])
164                (avltree::add t k less-than merge add) ) )
165
166;; Insert an element with the given KEY into TREE.  If an element with
167;; the same key is already in the tree, IF-FOUND is called with the
168;; element already there and it should return either it or the element
169;; to be inserted.  If an element with the same key is not in TREE,
170;; then IF-NOT-FOUND is called with no arguments and this should
171;; return the element to be inserted.
172;;-doc
173
174(define-inline (avltree::rotate-left! t)
175        (let ([r (%node-right t)])
176                (if (%empty-node? r)
177                        t
178                        (begin
179                                (%node-right-set! t (%node-left r))
180                                (%node-height-set! t (+ -1 (%node-height t)))
181                                (%node-left-set! r t)
182                                (%node-height-set! r (+ 1 (%node-height r)))
183                                r ) ) ) )
184
185(define-inline (avltree::rotate-right! t)
186        (let ([r (%node-left t)])
187                (if (%empty-node? r)
188                        t
189                        (begin
190                                (%node-left-set! t (%node-right r))
191                                (%node-height-set! t (+ -1 (%node-height t)))
192                                (%node-right-set! r t)
193                                (%node-height-set! r (+ 1 (%node-height r)))
194                                r ) ) ) )
195
196(define-inline (avltree::rotate! t)
197        (let ([lht (*node-height (%node-left t))]
198                                [rht (*node-height (%node-right t))])
199                (%node-height-set! t (avltree::max lht rht))
200                (cond
201                        [(> lht rht)
202                                (avltree::rotate-right! t)]
203                        [(< lht rht)
204                                (avltree::rotate-left! t)]
205                        [else
206                                t]) ) )
207
208(define (avltree::add! t k less-than if-found if-not-found)
209        (let loop ([t t])
210                (if (%empty-node? t)
211                        (if-not-found)
212                        (let ([ck (%node-key t)])
213                                (cond
214                                        [(less-than ck k)
215                                                (let ([n (loop (%node-right t))])
216                                                        (%node-right-set! t n)
217                                                        (avltree::rotate! t))]
218                                        [(less-than k ck)
219                                                (let ([n (loop (%node-left t))])
220                                                        (%node-left-set! t n)
221                                                        (avltree::rotate! t))]
222                                        [else
223                                                (if-found t)] ) ) ) ) )
224
225(define (avltree:merge-insert! t k less-than if-found if-not-found)
226        (let ([merge
227                                        (lambda (t)
228                                                (if (%node-deleted? t)
229                                                        (begin
230                                                                (%node-value-set! t (if-not-found))
231                                                                (%node-deleted-set! t #f))
232                                                        (%node-value-set! t (if-found (%node-value t))))
233                                                t)]
234                                [add
235                                        (lambda ()
236                                                (make-leaf k (if-not-found)))])
237                (avltree::add! t k less-than merge add) ) )
238
239;; Look for an element with the given KEY in TREE.  If a matching
240;; element is found IF-FOUND is called with the element.  If no
241;; matching element is found, IF-NOT-FOUND is called with no arguments.
242;;
243;; Marks the node!
244
245(define (avltree:delete t k less-than if-found if-not-found)
246        (let loop ([t t])
247                (if (%empty-node? t)
248                        (if-not-found)
249                        (let ([ck (%node-key t)])
250                                (cond
251                                        [(less-than k ck)
252                                                (loop (%node-left t))]
253                                        [(less-than ck k)
254                                                (loop (%node-right t))]
255                                        [(%node-deleted? t)
256                                                (if-not-found)]
257                                        [else
258                                                (if-found (%node-value t))
259                                                (%node-deleted-set! t #t)]) ) ) ) )
260
261;; Look for an element with the given KEY in TREE.  If a matching
262;; element is found IF-FOUND is called with the element.  If no
263;; matching element is found, IF-NOT-FOUND is called with no arguments.
264
265(define (avltree:find t k less-than if-found if-not-found)
266        (let loop ([t t])
267                (if (%empty-node? t)
268                        (if-not-found)
269                        (let ([ck (%node-key t)])
270                                (cond
271                                        [(less-than k ck)
272                                                (loop (%node-left t))]
273                                        [(less-than ck k)
274                                                (loop (%node-right t))]
275                                        [(%node-deleted? t)
276                                                (if-not-found)]
277                                        [else
278                                                (if-found (%node-value t))]) ) ) ) )
279
280;; Applies ACTION to each KEY DATA element in TREE in order, but requires
281;; NEXT to be called to drive the computation.
282
283(define (avltree:for-each-in-order t f s)
284        (let loop ([t t] [s s])
285                (cond
286                        [(%empty-node? t)
287                                s]
288                        [(%node-deleted? t)
289                                (loop (%node-right t) (loop (%node-left t) s))]
290                        [else
291                                (f (%node-key t) (%node-value t) (loop (%node-left t) s)
292                                        (lambda (s) (loop (%node-right t) s)))]) ) )
293
294;; Applies BEFORE and AFTER to all KEY DATA elements in TREE in order.
295;; BEFORE is applied during the downward pass and AFTER is applied on the
296;; upwards pass.
297
298(define (avltree:for-all-in-order t b a s)
299        (let loop ([t t] [s s])
300                (cond
301                        [(%empty-node? t)
302                                s]
303                        [(%node-deleted? t)
304                                (loop (%node-right t) (loop (%node-left t) s))]
305                        [else
306                                (let ([k (%node-key t)]
307                                                        [v (%node-value t)])
308                                        (loop (%node-right t) (a k v (loop (%node-left t) (b k v s)))))]) ) )
309
310;;
311
312(define (avltree:fold t func init)
313        (let loop ([t t] [r init])
314                (cond
315                        [(%empty-node? t)
316                                r]
317                        [(%node-deleted? t)
318                                (loop (%node-right t)
319                                        (loop (%node-left t) r))]
320                        [else
321                                (loop (%node-right t)
322                                        (func (%node-key t) (%node-value t)
323                                                (loop (%node-left t) r)))]) ) )
324
325;;
326
327(define (avltree:copy t less-than)
328        (avltree:fold t
329                (lambda (key val nt)
330                        (avltree:merge-insert nt key less-than (lambda (k v) v) (lambda () val)))
331                *empty-node*) )
332
333;;
334
335(define (avltree:size t)
336        (avltree:fold t (lambda (k v s) (+ 1 s)) 0) )
337
338;;
339
340(define (avltree:keys t)
341        (avltree:fold t (lambda (k v s) (cons k s)) '()) )
342
343;;
344
345(define (avltree:values t)
346        (avltree:fold t (lambda (k v s) (cons v s)) '()) )
347
348;;
349
350(define (avltree:to-alist t)
351        (avltree:fold t (lambda (k v s) (alist-cons k v s)) '()) )
352
353;;
354
355(define (avltree:from-alist l less-than if-dup)
356        (let loop ([t *empty-node*] [l l])
357                (cond
358                        [(null? l)
359                                t]
360                        [(pair? (car l))
361                                (let ([value (cdar l)])
362                                        (loop
363                                                (avltree:merge-insert t (caar l) less-than
364                                                        (lambda (o) (if-dup o value)) (lambda () value))
365                                                (cdr l)))]
366                        [else
367                                (let ([value (car l)])
368                                        (loop
369                                                (avltree:merge-insert t value less-than
370                                                        (lambda (o) (if-dup o value)) (lambda () value))
371                                                (cdr l)))
372                                #;(error 'alist->avltree "invalid association list" l)]) ) )
373
374;;
375
376#|
377(define-inline (spaces out n)
378        (let loop ([n n])
379                (unless (zero? n)
380                        (display #\space out)
381                        (loop (sub1 n))) ) )
382
383(define (avltree:print out t indent)
384        (if (%empty-node? t)
385                (begin
386                        (spaces out indent) (fprintf out "()~%"))
387                (begin
388                        (spaces out indent) (fprintf out "~A~%" t)
389                        (avltree:print out (%node-left t) (+ 2 indent))
390                        (avltree:print out (%node-right t) (+ 2 indent)))) )
391|#
392
393(define (avltree:print nod out)
394        (if (%empty-node? nod)
395                (write '() out)
396                (begin
397                  (display #\( out)
398      (write nod out)
399                  (display #\space out)
400      (avltree:print (%node-left nod) out)
401                  (display #\space out)
402      (avltree:print (%node-right nod) out)
403      (display #\) out))) )
404
405;;;
406
407(define-inline-unchecked-record-type avltree
408        (%make-tree root less-than)
409        %tree?
410        (root %tree-root %tree-root-set!)
411        (less-than %tree-less-than) )
412
413(define-record-printer (avltree tree out)
414        (avltree:print (%tree-root tree) out) )
415
416(define-inline (check-tree obj loc)
417        (unless (%tree? obj)
418                (error loc "invalid avltree" obj)) )
419
420(define-inline (check-list obj loc)
421        (unless (list? obj)
422                (error loc "invalid list" obj)) )
423
424(define-inline (check-procedure obj loc)
425        (unless (procedure? obj)
426                (error loc "invalid procedure" obj)) )
427
428(define (not-found-error)
429        (error 'avltree "not found") )
430
431;;;
432
433(define (make-avltree less-than)
434        (check-procedure less-than 'make-avltree)
435        (%make-tree *empty-node* less-than) )
436
437(define (alist->avltree lst less-than #!optional (if-dup (lambda (o n) n)))
438        (check-list lst 'alist->avltree)
439        (check-procedure less-than 'alist->avltree)
440        (check-procedure if-dup 'alist->avltree)
441        (%make-tree (avltree:from-alist lst less-than if-dup) less-than) )
442
443(define (avltree? obj)
444        (%tree? obj) )
445
446(define (avltree-empty? tree)
447        (check-tree tree 'avltree-empty?)
448        (%empty-node? (%tree-root tree)) )
449
450(define (avltree-less-than-function tree)
451        (check-tree tree 'avltree-empty?)
452        (%tree-less-than tree) )
453
454(define (avltree-size tree)
455        (check-tree tree 'avltree-size)
456        (avltree:size (%tree-root tree)) )
457
458(define (avltree-keys tree)
459        (check-tree tree 'avltree-keys)
460        (avltree:keys (%tree-root tree)) )
461
462(define (avltree-values tree)
463        (check-tree tree 'avltree-values)
464        (avltree:values (%tree-root tree)) )
465
466(define (avltree->alist tree)
467        (check-tree tree 'avltree->alist)
468        (avltree:to-alist (%tree-root tree)) )
469
470(define (avltree-ref tree key #!optional (if-not-found not-found-error))
471        (check-tree tree 'avltree-find)
472        (check-procedure if-not-found 'avltree-ref)
473        (avltree:find (%tree-root tree) key (%tree-less-than tree)
474                identity if-not-found) )
475
476(define (avltree-ref/default tree key def)
477        (check-tree tree 'avltree-find)
478        (avltree:find (%tree-root tree) key (%tree-less-than tree)
479                identity (lambda () def)) )
480
481(define (avltree-exists? tree key)
482        (check-tree tree 'avltree-exists?)
483        (avltree:find (%tree-root tree) key (%tree-less-than tree)
484                true false) )
485
486(define (avltree-update! tree key if-found #!optional (if-not-found not-found-error))
487        (check-tree tree 'avltree-update!)
488        (check-procedure if-found 'avltree-update!)
489        (check-procedure if-not-found 'avltree-update!)
490        (%tree-root-set! tree
491                (avltree:merge-insert! (%tree-root tree) key (%tree-less-than tree)
492                        if-found if-not-found)) )
493
494(define (avltree-update!/default tree key if-found def)
495        (check-tree tree 'avltree-update!/default)
496        (check-procedure if-found 'avltree-update!/default)
497        (%tree-root-set! tree
498                (avltree:merge-insert! (%tree-root tree) key (%tree-less-than tree)
499                        if-found (lambda () def))) )
500
501(define (avltree-set! tree key value #!optional (if-found (lambda (v) value)))
502        (check-tree tree 'avltree-set!)
503        (check-procedure if-found 'avltree-set!)
504        (%tree-root-set! tree
505                (avltree:merge-insert! (%tree-root tree) key (%tree-less-than tree)
506                        if-found (lambda () value))) )
507
508(define (avltree-delete! tree key #!optional (if-found identity) (if-not-found false))
509        (check-tree tree 'avltree-delete!)
510        (check-procedure if-found 'avltree-delete!)
511        (check-procedure if-not-found 'avltree-delete!)
512        (avltree:delete (%tree-root tree) key (%tree-less-than tree)
513                if-found if-not-found) )
514
515(define (avltree-vacuum! tree)
516        (check-tree tree 'avltree-vacuum!)
517        (%tree-root-set! tree (avltree:copy (%tree-root tree) (%tree-less-than tree))) )
518
519(define (avltree-merge! tree1 tree2)
520        (check-tree tree1 'avltree-merge!)
521        (check-tree tree2 'avltree-merge!)
522        (let ([less-than (%tree-less-than tree1)])
523                (%tree-root-set! tree1
524                        (avltree:for-each-in-order (%tree-root tree2)
525                                (lambda (nd key val nxt)
526                                        (nxt (avltree:merge-insert! nd key less-than (lambda (v) val) (lambda () val))))
527                                (%tree-root tree1))) ) )
528
529(define (avltree-update tree key if-found #!optional (if-not-found not-found-error))
530        (check-tree tree 'avltree-update)
531        (check-procedure if-found 'avltree-update)
532        (check-procedure if-not-found 'avltree-update)
533        (%make-tree
534                (avltree:merge-insert (%tree-root tree) key (%tree-less-than tree)
535                        if-found if-not-found)
536                (%tree-less-than tree)) )
537
538(define (avltree-update/default tree key if-found def)
539        (check-tree tree 'avltree-update/default)
540        (check-procedure if-found 'avltree-update/default)
541        (%make-tree
542                (avltree:merge-insert (%tree-root tree) key (%tree-less-than tree)
543                        if-found (lambda () def))
544                (%tree-less-than tree)) )
545
546(define (avltree-set tree key value #!optional (if-found (lambda (v) value)))
547        (check-tree tree 'avltree-set)
548        (check-procedure if-found 'avltree-set)
549        (%make-tree
550                (avltree:merge-insert (%tree-root tree) key (%tree-less-than tree)
551                        if-found (lambda () value))
552                (%tree-less-than tree)) )
553
554(define (avltree-delete tree key #!optional (if-found identity) (if-not-found false))
555        (check-tree tree 'avltree-delete)
556        (check-procedure if-found 'avltree-delete)
557        (check-procedure if-not-found 'avltree-delete)
558        (let* ([less-than (%tree-less-than tree)]
559                                 [ntree (%make-tree (avltree:copy (%tree-root tree) less-than) less-than)])
560                (avltree:delete (%tree-root ntree) key less-than
561                        if-found if-not-found)
562                ntree ) )
563
564(define (avltree-copy tree)
565        (check-tree tree 'avltree-copy)
566        (let ([less-than (%tree-less-than tree)])
567                (%make-tree (avltree:copy (%tree-root tree) less-than) less-than) ) )
568
569(define (avltree-merge tree1 tree2)
570        (check-tree tree1 'avltree-merge)
571        (check-tree tree2 'avltree-merge)
572        (let* ([less-than (%tree-less-than tree1)]
573                                 [ntree (%make-tree (avltree:copy (%tree-root tree1) less-than) less-than)])
574                (%tree-root-set! ntree
575                        (avltree:for-each-in-order (%tree-root tree2)
576                                (lambda (nd key val nxt)
577                                        (nxt (avltree:merge-insert nd key less-than (lambda (v) val) (lambda () val))))
578                                (%tree-root ntree)))
579                ntree ) )
580
581(define (avltree-walk tree proc)
582        (check-tree tree 'avltree-walk)
583        (avltree:fold (%tree-root tree) (lambda (k v s) (proc k v)) (void)) )
584
585(define (avltree-fold tree func init)
586        (check-tree tree 'avltree-fold)
587        (avltree:fold (%tree-root tree) func init) )
588
589(define (avltree-enfold tree func init)
590        (check-tree tree 'avltree-enfold)
591        (avltree:for-each-in-order (%tree-root tree) func init) )
592
593(define (avltree-bifold tree before after init)
594        (check-tree tree 'avltree-bifold)
595        (avltree:for-all-in-order (%tree-root tree) before after init) )
Note: See TracBrowser for help on using the repository browser.