source: project/release/3/rb-tree/trunk/rb-tree.scm @ 12429

Last change on this file since 12429 was 12429, checked in by Ivan Raikov, 11 years ago

Fixes to for-ascending and descending.

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