source: project/rb-tree/trunk/rb-tree.scm @ 8037

Last change on this file since 8037 was 8037, checked in by Ivan Raikov, 12 years ago

Replaced values and let-values with list and match-let.

File size: 19.1 KB
Line 
1;;
2;; Red-black tree
3;;
4;; An implementation of an ordered dictionary data structure, based
5;; on red-black trees.
6;;
7;; This code is based on the SML/NJ library implementation of
8;; red-black trees, which is in turn based on Chris Okasaki's
9;; implementation of red-black trees.  The delete function is based on
10;; the description in Cormen, Leiserson, and Rivest.
11;;
12;; Some helper code was borrowed from treap.scm by Oleg Kiselyov.
13;;
14;;
15;; Copyright 2007 Ivan Raikov and the Okinawa Institute of Science and Technology
16;;
17;;
18;; This program is free software: you can redistribute it and/or
19;; modify it under the terms of the GNU General Public License as
20;; published by the Free Software Foundation, either version 3 of the
21;; License, or (at your option) any later version.
22;;
23;; This program is distributed in the hope that it will be useful, but
24;; WITHOUT ANY WARRANTY; without even the implied warranty of
25;; MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
26;; General Public License for more details.
27;;
28;; A full copy of the GPL license can be found at
29;; <http://www.gnu.org/licenses/>.
30;;
31;;
32;; TODO: Add the linear-time tree construction code from the
33;; paper _Constructing red-black trees_ by Hinze.
34;;
35
36(require-extension srfi-1)
37(require-extension match)
38(require-extension datatype)
39
40(define-extension rb-tree)
41
42(declare (export  make-rb-tree))
43
44
45;;
46;; A red-black tree should satisfy the following two invariants:
47;;
48;;   Red Invariant: each red node has a black parent.
49;;
50;;   Black Condition: each path from the root to an empty node has the
51;;     same number of black nodes (the tree's black height).
52;;
53;; The Red condition implies that the root is always black and the Black
54;; condition implies that any node with only one child will be black and
55;; its child will be a red leaf.
56;;
57;;
58;; The red-black tree object is created by procedure make-rb-tree, the
59;; only user-visible function defined in this library:
60;;
61;;  make-rb-tree:: KEY-COMPARE-PROC -> RB-TREE
62;;
63;;  where KEY-COMPARE-PROC is a user-supplied function
64;;
65;;  KEY-COMPARE-PROC:: key1 key2 -> INTEGER
66;;
67;;  that takes two keys and returns a negative, positive, or zero
68;;  number depending on how the first key compares to the second.
69;;
70;; The red-black tree object responds to the following messages
71;;
72;;      'get
73;;
74;;              returns a procedure LAMBDA KEY . DEFAULT-CLAUSE which
75;;              searches the red-black tree for an association with a
76;;              given KEY, and returns a (key . value) pair of the
77;;              found association.  If an association with the KEY
78;;              cannot be located in the red-black tree, the PROC
79;;              returns the result of evaluating the DEFAULT-CLAUSE.
80;;              If the default clause is omitted, an error is
81;;              signalled.  The KEY must be comparable to the keys in
82;;              the red-black tree by a key-compare predicate (which
83;;              has been specified when the red-black tree was
84;;              created)
85;;
86;;      'get-min
87;;
88;;              returns a (key . value) pair for an association in the
89;;              red-black tree with the smallest key. If the red-black
90;;              tree is empty, an error is signalled.
91;;
92;;      'delete-min!
93;;
94;;              removes the min key and the corresponding association
95;;              from the red-black tree. Returns a (key . value) pair
96;;              of the removed association.  If the red-black tree is
97;;              empty, an error is signalled.
98;;
99;;      'get-max
100;;
101;;              returns a (key . value) pair for an association in the
102;;              red-black tree with the largest key. If the red-black
103;;              tree is empty, an error is signalled.
104;;
105;;      'delete-max!
106;;
107;;              removes the max key and the corresponding association
108;;              from the red-black tree. Returns a (key . value) pair
109;;              of the removed association.  If the red-black tree is
110;;              empty, an error is signalled.
111;;
112;;      empty?
113;;              returns #t if the red-black tree is empty
114;;
115;;      size
116;;
117;;              returns the size (the number of associations) in the
118;;              red-black tree
119;;
120;;      depth
121;;
122;;              returns the depth of the tree. It requires the
123;;              complete traversal of the tree, so use sparingly
124;;
125;;      clear!
126;;
127;;              removes all associations from the red-black tree (thus
128;;              making it empty)
129;;
130;;      'put!
131;;
132;;              returns a procedure LAMBDA KEY VALUE which, given a
133;;              KEY and a VALUE, adds the corresponding association to
134;;              the red-black tree. If an association with the same
135;;              KEY already exists, its value is replaced with the
136;;              VALUE (and the old (key . value) association is
137;;              returned). Otherwise, the return value is #f.
138;;
139;;      'delete!
140;;
141;;              returns a procedure LAMBDA KEY . DEFAULT-CLAUSE which
142;;              searches the red-black tree for an association with a
143;;              given KEY, deletes it, and returns a (key . value)
144;;              pair of the found and deleted association.  If an
145;;              association with the KEY cannot be located in the
146;;              red-black tree, the PROC returns the result of
147;;              evaluating the DEFAULT-CLAUSE.  If the default clause
148;;              is omitted, an error is signalled.
149;;
150;;      for-each-ascending
151;;
152;;              returns a procedure LAMBDA PROC that will apply the
153;;              given procedure PROC to each (key . value) association
154;;              of the red-black tree, from the one with the smallest
155;;              key all the way to the one with the max key, in an
156;;              ascending order of keys.  The red-black tree must not
157;;              be empty.
158;;
159;;      for-each-descending
160;;
161;;              returns a procedure LAMBDA PROC that will apply the
162;;              given procedure PROC to each (key . value) association
163;;              of the red-black tree, in the descending order of
164;;              keys.  The red-black tree must not be empty.
165;;
166
167
168(define (rb-tree:error x . rest)
169  (let ((port (open-output-string)))
170    (let loop ((objs (cons x rest)))
171      (if (null? objs)
172          (begin
173            (newline port)
174            (error 'rb-tree (get-output-string port)))
175          (begin (display (car objs) port)
176                 (display " " port)
177                 (loop (cdr objs)))))))
178
179(define R 'Red)
180(define B 'Black)
181
182(define (color? x) (or (eq? x 'Red) (eq? x 'Black)))
183
184(define-datatype tree tree?
185  (Empty)
186  (Tree  (color color?) (left tree?) (key identity) (value identity) (right tree?)))
187
188(define-datatype zipper zipper?
189  (Top) 
190  (Left (color color?) (key identity) (value identity) (tree tree?) (zipper zipper?))
191  (Right (color color?) (tree tree?) (key identity) (value identity) (zipper zipper?)))
192
193
194(define (make-rb-tree key-compare)
195  ;;
196  ;; This macro was borrowed from treap.scm by Oleg Kiselyov
197  ;;
198  (define-macro (dispatch-on-key key node-key on-less on-equal on-greater)
199    (let ((result (gensym)))
200      `(let ((,result (key-compare ,node-key ,key )))
201        (cond
202          ((zero? ,result) ,on-equal)
203          ((positive? ,result) ,on-greater)
204          (else ,on-less)))))
205
206  (let ((root (Empty)) (size 0))
207
208    (define (make-rb-tree-dispatcher root size)
209
210    ;; Adds a new association to the tree (or replaces the old one if
211    ;; existed). Returns the (key . value) pair of the old
212    ;; association, or #f if a new association was really added
213    (define (insert root key value)
214      (define (ins root)
215        (cases tree root
216               (Empty ()  (list #f (Tree R (Empty) key value (Empty))))
217               (Tree (color a yk y b) 
218                     (dispatch-on-key 
219                      key yk 
220                      ;; Case 1: key < yk
221                      (match a
222                             (($ tree 'Tree 'Red c zk z d)
223                              (dispatch-on-key 
224                               key zk
225                               ;; Case 1.1: key < zk
226                               (match-let (((found? c1) (ins c)))
227                                          (list found?
228                                                (match c1
229                                                       (($ tree 'Tree 'Red e wk w f)
230                                                        (Tree R (Tree B e wk w f) zk z (Tree B d yk y b)))
231                                                       (else  (Tree B (Tree R c1 zk z d) yk y b)))))
232                               ;; Case 1.2: key = zk
233                               (list a (Tree color (Tree R c key value d) yk y b))
234                               ;; Case 1.3: key > zk
235                               (match-let (((found? d1) (ins d)))
236                                          (list found?
237                                                (match d1
238                                                       (($ tree 'Tree 'Red e wk w f)
239                                                        (Tree R (Tree B c zk z e)  wk  w  (Tree B f yk y b)))
240                                                       (else (Tree B (Tree R c zk z d1) yk y b)))))))
241                             (else  (match-let (((found? a1)  (ins a)))
242                                               (list found? (Tree B a1 yk y b)))))
243                      ;; Case 2: key  = yk
244                      (list root (Tree color a key value b))
245                      ;; Case 3: key  > yk
246                      (match b
247                             (($ tree 'Tree 'Red c zk z d)
248                              (dispatch-on-key 
249                               key zk
250                               ;; Case 3.1: key < zk
251                               (match-let (((found? c1) (ins c)))
252                                          (list found?
253                                                (match c1
254                                                       (($ tree 'Tree 'Red e wk w f)
255                                                        (Tree R (Tree B a yk y e)  wk  w (Tree B f zk z d)))
256                                                       (else (Tree B a yk y (Tree R c1 zk z d))))))
257                               ;; Case 3.2: key = zk
258                               (list b (Tree color a yk y (Tree R c key value d)))
259                               ;; Case 3.3: key > zk
260                               (match-let (((found? d1) (ins d)))
261                                          (list found?
262                                                (match d1
263                                                       (($ tree 'Tree 'Red e wk w f)
264                                                        (Tree R (Tree B a yk y c)  zk z (Tree B e wk w f)))
265                                                       (else (Tree B a yk y (Tree R c zk z d1))))))))
266                             (else (match-let (((found? b1) (ins b)))
267                                               (list found? (Tree B a yk y b1)))))))))
268
269      (ins root))
270
271
272    ;; Looks for an item: Given a key, returns the corresponding (key
273    ;; . value) association or #f if the tree does not contain an
274    ;; association with that key.
275    (define (find-assoc key)
276      (define (find root)
277        (cases tree root 
278               (Empty ()  #f)
279               (Tree (c a yk y b)
280                     (dispatch-on-key 
281                      key yk (find a) (cons yk y) (find b)))))
282      (find root))
283
284    ;; Finds an association with a given key, and deletes it.  Returns
285    ;; the (key . value) pair of the deleted association, or #f if it
286    ;; couldn't be found
287    (define (delete root key)
288
289      (define (zip zipper tree)
290        (match (cons zipper tree)
291               ((($ zipper 'Top) . a)  tree)
292               ((($ zipper 'Left color xk x b z) . a)   (zip z (Tree color a xk x b)))
293               ((($ zipper 'Right color a xk x z) . b)  (zip z (Tree color a xk x b)))))
294
295      ;; bbZip propagates a black deficit up the tree until either
296      ;; the top is reached, or the deficit can be covered.  It
297      ;; returns a boolean that is true if there is still a deficit
298      ;; and the zipped tree.
299     (define (bbZip zipper tree)
300        (match (cons zipper tree)
301               ((($ zipper 'Top) . a)  (cons #t a))
302               ;; case 1L
303               ((($ zipper 'Left 'Black xk x ($ tree 'Tree 'Red c yk y d) z) . a)
304                (bbZip (Left R xk x c (Left B yk y d z)) a))
305               ;; case 3L
306               ((($ zipper 'Left color xk x ($ tree 'Tree 'Black ($ tree 'Tree 'Red c yk y d) wk w e) z) . a)
307                (bbZip (Left color xk x (Tree B c yk y (Tree R d wk w e)) z) a))
308               ;; case 4L
309               ((($ zipper 'Left color xk x ($ tree 'Tree 'Black c yk y ($ tree 'Tree 'Red d wk w e)) z) . a)
310                (cons #f (zip z (Tree color (Tree B a xk x c) yk y (Tree B d wk w e)))))
311               ;; case 2L
312               ((($ zipper 'Left 'Red xk x ($ tree 'Tree 'Black c yk y d) z) . a)
313                (cons #f (zip z (Tree B a xk x (Tree R c yk y d)))))
314               ;; case 2L
315               ((($ zipper 'Left 'Black xk x ($ tree 'Tree 'Black c yk y d) z) . a)
316                (bbZip z (Tree B a xk x (Tree R c yk y d))))
317               ;; case 1R
318               ((($ zipper 'Right color ($ tree 'Tree 'Red c yk y d) xk x z) . b) 
319                (bbZip (Right R d xk x (Right B c yk y z)) b))
320               ;; case 3R
321               ((($ zipper 'Right color ($ tree 'Tree 'Black ($ tree 'Tree 'Red c wk w d) yk y e) xk x z) . b) 
322                (bbZip (Right color (Tree B c wk w (Tree R d yk y e)) xk x z) b))
323               ;; case 4R
324               ((($ zipper 'Right color ($ tree 'Tree 'Black c yk y ($ tree 'Tree 'Red d wk w e)) xk x z) . b) 
325                (cons #f (zip z (Tree color c yk y (Tree B (Tree R d wk w e) xk x b)))))
326               ;; case 2R
327               ((($ zipper 'Right 'Red ($ tree 'Tree 'Black c yk y d) xk x z) . b) 
328                (cons #f (zip z (Tree B (Tree R c yk y d) xk x b))))
329               ;; case 2R
330               ((($ zipper 'Right 'Black ($ tree 'Tree 'Black c yk y d) xk x z) . b) 
331                (bbZip z (Tree B (Tree R c yk y d) xk x b)))
332               (else   (cons #f (zip zipper tree)))))
333
334      (define (delMin tree z)
335        (match tree
336               (($ tree 'Tree 'Red ($ tree 'Empty) yk y b) 
337                (list yk y (cons #f (zip z b))))
338               (($ tree 'Tree 'Black ($ tree Empty) yk y b) 
339                (list yk y (bbZip z b)))
340               (($ tree 'Tree color a yk y b) 
341                (delMin a (Left color yk y b z)))
342               (($ tree 'Empty) (rb-tree:error 'delete! "invalid tree"))))
343
344      (define (join color a b z)
345        (match (list color a b)
346               (( 'Red ($ tree 'Empty) ($ tree 'Empty)) 
347                (zip z (Empty)))
348               (( _ a ($ tree 'Empty)) 
349                (cdr  (bbZip z a)))
350               (( _ ($ tree 'Empty) b)
351                (cdr  (bbZip z b)))
352               (( color a b)
353                (match-let (((xk x b)  (delMin b (Top))))
354                           (match b
355                                  ((#t . b1)  (cdr  (bbZip z (Tree color a xk x b1))))
356                                  ((#f . b1)  (zip z (Tree color a xk x b1))))))))
357
358
359      (define (del tree z)
360        (match tree 
361               (($ tree 'Empty)  #f)
362               (($ tree 'Tree color a yk y b) 
363                (dispatch-on-key 
364                 key yk 
365                 (del a (Left color yk y b z))
366                 (cons (cons yk y) (join color a b z))
367                 (del b (Right color a yk y z))))))
368
369      (del root (Top)))
370
371    (define (delete! key)
372      (let ((item+tree  (delete root key)))
373        (and item+tree
374             (begin
375               (set! root (cdr item+tree))
376               (set! size (- size 1))
377               (car item+tree)))))
378
379
380    (define (get-min)
381      (define (f root)
382        (match root
383               (($ tree 'Empty)  #f)
384               (($ tree 'Tree _ _ ($ tree 'Empty) xk x _)  (cons xk x))
385               (($ tree 'Tree _ a _ _ _)   (f a))))
386      (f root))
387
388    (define (get-max)
389      (define (f root)
390        (match root
391               (($ tree 'Empty)  #f)
392               (($ tree 'Tree _ _ xk x ($ tree 'Empty))  (cons xk x))
393               (($ tree 'Tree _ _ _ _ b)   (f b))))
394      (f root))
395
396
397    (define (fold-limit p f init) 
398      (define (foldf tree ax)
399        (match tree
400               (($ tree 'Empty)  ax)
401               (($ tree 'Tree _ a _ x b) 
402                (if (p ax) ax (foldf b (f x (foldf a ax)))))))
403      (foldf root init))
404
405    (define (fold-right-limit p f init) 
406      (define (foldf tree ax)
407        (match tree
408               (($ tree 'Empty)  ax)
409               (($ tree 'Tree _ a _ x b) 
410                (if (p ax) ax (foldf a (f x (foldf b ax)))))))
411      (foldf root init))
412
413    (define (fold-partial p f init) 
414      (define (foldf tree ax)
415        (match tree
416               (($ tree 'Empty)  ax)
417               (($ tree 'Tree _ a _ x b) 
418                (if (p x) (foldf b (f x (foldf a ax))) ax))))
419      (foldf root init))
420
421    (define (foldi-partial p f init) 
422      (define (foldf tree ax)
423        (match tree
424               (($ tree 'Empty)  ax)
425               (($ tree 'Tree _ a xk x b) 
426                (if (p xk x) (foldf b (f xk x (foldf a ax))) ax))))
427      (foldf root init))
428
429    (define (fold-right-partial p f init) 
430      (define (foldf tree ax)
431        (match tree
432               (($ tree 'Empty)  ax)
433               (($ tree 'Tree _ a _ x b) 
434                (if (p x) (foldf a (f x (foldf b ax))) ax))))
435      (foldf root init))
436
437    (define (foldi-right-partial p f init) 
438      (define (foldf tree ax)
439        (match tree
440               (($ tree 'Empty)  ax)
441               (($ tree 'Tree _ a xk x b) 
442                (if (p xk x) (foldf a (f xk x (foldf b ax))) ax))))
443      (foldf root init))
444
445
446    (define (fold f init)
447      (define (foldf tree ax)
448        (match tree
449               (($ tree 'Empty)  ax)
450               (($ tree 'Tree _ a _ x b)  (foldf b (f x (foldf a ax))))))
451      (foldf root init))
452
453    (define (foldi f init)
454      (define (foldf tree ax)
455        (match tree
456               (($ tree 'Empty)  ax)
457               (($ tree 'Tree _ a xk x b)  (foldf b (f xk x (foldf a ax))))))
458      (foldf root init))
459
460    (define (fold-right f init)
461      (define (foldf tree ax)
462        (match tree
463               (($ tree 'Empty)  ax)
464               (($ tree 'Tree _ a _ x b)  (foldf a (f x (foldf b ax))))))
465      (foldf root init))
466
467    (define (foldi-right f init)
468      (define (foldf tree ax)
469        (match tree
470               (($ tree 'Empty)  ax)
471               (($ tree 'Tree _ a xk x b)  (foldf a (f xk x (foldf b ax))))))
472      (foldf root init))
473
474
475    (define (get-depth)
476      (let loop ((node root) (level 0))
477        (match node 
478               (($ tree 'Empty)  level)
479               (($ tree 'Tree _ a _ _ b)  (max (loop a (+ 1 level))
480                                                (loop b (+ 1 level)))))))
481
482    ;; Returns an ordered list of the keys in the tree
483    (define (list-keys)
484      (foldi-right (lambda (k x l) (cons k l)) (list)))
485
486    ;; Returns an ordered list of the (key . item) pairs in the tree
487    (define (list-items)
488      (foldi-right (lambda (k x l) (cons (cons k x) l)) (list)))
489
490    (define (for-each-ascending f)
491      (define (appf tree)
492        (match tree
493               (($ tree 'Empty)  (void))
494               (($ tree 'Tree _ a _ x b)  (begin (appf a) (f x) (appf b)))))
495      (appf root))
496
497    (define (for-each-descending f)
498      (define (appf tree)
499        (match tree
500               (($ tree 'Empty)  (void))
501               (($ tree 'Tree _ a _ x b)  (begin (appf b) (f x) (appf a)))))
502      (appf root))
503
504    (define (map f)
505      (define (mapf tree)
506        (match tree
507               (($ tree 'Empty)  (Empty))
508               (($ tree 'Tree color a xk x b) 
509                (Tree color (mapf a) xk (f x) (mapf b)))))
510      (make-rb-tree-dispatcher (mapf root) size))
511
512    (define (mapi f)
513      (define (mapf tree)
514        (match tree
515               (($ tree 'Empty)   (Empty))
516               (($ tree 'Tree color a xk x b) 
517                 (Tree color (mapf a) xk (f xk x) (mapf b)))))
518      (make-rb-tree-dispatcher  (mapf root) size))
519
520
521    (define (apply-default-clause label key default-clause)
522      (cond
523        ((null? default-clause)
524          (rb-tree:error label "key " key " was not found in the tree"))
525        ((pair? (cdr default-clause))
526          (rb-tree:error label "default argument must be a single clause"))
527        ((procedure? (car default-clause)) ((car default-clause)))
528        (else (car default-clause))))
529   
530    ;; Dispatcher
531    (lambda (selector)
532      (case selector
533        ((get)
534          (lambda (key . default-clause)
535            (or (find-assoc key) (apply-default-clause 'get key default-clause))))
536
537        ((delete!)
538          (lambda (key . default-clause)
539            (or (delete! key)
540                (apply-default-clause 'delete! key default-clause))))
541
542        ((delete)
543          (lambda (key . default-clause)
544            (or (let ((item+tree  (delete root key)))
545                  (and item+tree
546                       (make-rb-tree-dispatcher  (cdr item+tree) 
547                                                 (if (car item+tree) (- size 1) size))))
548                (apply-default-clause 'delete key default-clause))))
549
550        ((put!) 
551         (lambda (key value)
552           (match-let (((found? new-root)  (insert root key value)))
553                      (set! root new-root)
554                      (if (not found?)  (set! size (+ 1 size)))
555                      found?)))
556
557        ((put) 
558         (lambda (key value)
559           (match-let (((found? new-root)  (insert root key value)))
560                      (make-rb-tree-dispatcher  new-root (if (not found?) (+ 1 size) size)))))
561
562
563        ((get-min) (get-min))
564        ((get-max) (get-max))
565
566        ((delete-min!)   (delete! (car (get-min))))
567        ((delete-max!)   (delete! (car (get-max))))
568
569        ((empty?)  (cases tree root
570                          (Empty () #t)
571                          (else #f)))
572        ((size)    size)
573
574        ((depth)   (get-depth))
575
576        ((clear!)  (begin
577                     (set! root (Empty)) 
578                     (set! size 0)))
579
580        ((for-each-ascending)   for-each-ascending)
581        ((for-each-descending)  for-each-descending)
582
583        ((list-keys)            (list-keys))
584        ((list-items)           (list-items))
585       
586        ((map)                  map)
587        ((mapi)                 mapi)
588
589        ((fold)                fold)
590        ((foldi)               foldi)
591        ((fold-right)          fold-right)
592        ((foldi-right)         foldi-right)
593
594        ((fold-partial)        fold-partial)
595        ((foldi-partial)       foldi-partial)
596        ((fold-right-partial)          fold-right-partial)
597        ((foldi-right-partial)         foldi-right-partial)
598
599        ((fold-limit)          fold-limit)
600        ((fold-right-limit)    fold-right-limit)
601
602        (else
603          (rb-tree:error 'selector "unknown message " selector " sent to a red-black tree")))))
604
605  (make-rb-tree-dispatcher root size)))
606
Note: See TracBrowser for help on using the repository browser.