home *** CD-ROM | disk | FTP | other *** search
/ OS/2 Shareware BBS: 10 Tools / 10-Tools.zip / mitsch75.zip / scheme-7_5_17-src.zip / scheme-7.5.17 / src / runtime / rbtree.scm < prev    next >
Text File  |  2000-02-14  |  14KB  |  465 lines
  1. #| -*-Scheme-*-
  2.  
  3. $Id: rbtree.scm,v 1.7 2000/02/14 19:59:44 cph Exp $
  4.  
  5. Copyright (c) 1993-2000 Massachusetts Institute of Technology
  6.  
  7. This program is free software; you can redistribute it and/or modify
  8. it under the terms of the GNU General Public License as published by
  9. the Free Software Foundation; either version 2 of the License, or (at
  10. your option) any later version.
  11.  
  12. This program is distributed in the hope that it will be useful, but
  13. WITHOUT ANY WARRANTY; without even the implied warranty of
  14. MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
  15. General Public License for more details.
  16.  
  17. You should have received a copy of the GNU General Public License
  18. along with this program; if not, write to the Free Software
  19. Foundation, Inc., 675 Mass Ave, Cambridge, MA 02139, USA.
  20. |#
  21.  
  22. ;;;; Red-Black Trees
  23. ;;; package: (runtime rb-tree)
  24.  
  25. ;;; Cormen, Leiserson, and Rivest, "Introduction to Algorithms",
  26. ;;; Chapter 14, "Red-Black Trees".
  27.  
  28. ;;; Properties of Red-Black Trees:
  29. ;;; 1. Every node is either red or black.
  30. ;;; 2. Every leaf (#F) is black.
  31. ;;; 3. If a node is red, then both its children are black.
  32. ;;; 4. Every simple path from a node to a descendent leaf contains the
  33. ;;;    same number of black nodes.
  34. ;;; These algorithms additionally assume:
  35. ;;; 5. The root of a tree is black.
  36.  
  37. (declare (usual-integrations))
  38.  
  39. (define-structure (tree
  40.            (predicate rb-tree?)
  41.            (constructor make-tree (key=? key<?)))
  42.   (root #f)
  43.   (key=? #f read-only #t)
  44.   (key<? #f read-only #t))
  45.  
  46. (define (make-rb-tree key=? key<?)
  47.   ;; Optimizations to work around compiler that codes known calls to
  48.   ;; these primitives much more efficiently than unknown calls.
  49.   (make-tree (cond ((eq? key=? eq?) (lambda (x y) (eq? x y)))
  50.            ((eq? key=? fix:=) (lambda (x y) (fix:= x y)))
  51.            ((eq? key=? flo:=) (lambda (x y) (flo:= x y)))
  52.            (else key=?))
  53.          (cond ((eq? key<? fix:<) (lambda (x y) (fix:< x y)))
  54.            ((eq? key<? flo:<) (lambda (x y) (flo:< x y)))
  55.            (else key<?))))
  56.  
  57. (define-integrable (guarantee-rb-tree tree procedure)
  58.   (if (not (rb-tree? tree))
  59.       (error:wrong-type-argument tree "red/black tree" procedure)))
  60.  
  61. (define-structure (node
  62.            (constructor make-node (key datum)))
  63.   key
  64.   datum
  65.   (up #f)
  66.   (left #f)
  67.   (right #f)
  68.   (color #f))
  69.  
  70. ;;; The algorithms are left/right symmetric, so abstract "directions"
  71. ;;; permit code to be used for either symmetry:
  72.  
  73. (define-integrable (b->d left?)
  74.   (if left? 'LEFT 'RIGHT))
  75.  
  76. (define-integrable (-d d)
  77.   (if (eq? 'LEFT d) 'RIGHT 'LEFT))
  78.  
  79. (define-integrable (get-link+ p d)
  80.   (if (eq? 'LEFT d)
  81.       (node-left p)
  82.       (node-right p)))
  83.  
  84. (define-integrable (set-link+! p d l)
  85.   (if (eq? 'LEFT d)
  86.       (set-node-left! p l)
  87.       (set-node-right! p l)))
  88.  
  89. (define-integrable (get-link- p d)
  90.   (if (eq? 'RIGHT d)
  91.       (node-left p)
  92.       (node-right p)))
  93.  
  94. (define-integrable (set-link-! p d l)
  95.   (if (eq? 'RIGHT d)
  96.       (set-node-left! p l)
  97.       (set-node-right! p l)))
  98.  
  99. (define (rotate+! tree x d)
  100.   ;; Assumes (NOT (NOT (GET-LINK- X D))).
  101.   (let ((y (get-link- x d)))
  102.     (let ((beta (get-link+ y d)))
  103.       (set-link-! x d beta)
  104.       (if beta (set-node-up! beta x)))
  105.     (let ((u (node-up x)))
  106.       (set-node-up! y u)
  107.       (cond ((not u)
  108.          (set-tree-root! tree y))
  109.         ((eq? x (get-link+ u d))
  110.          (set-link+! u d y))
  111.         (else
  112.          (set-link-! u d y))))
  113.     (set-link+! y d x)
  114.     (set-node-up! x y)))
  115.  
  116. (define-integrable (rotate-! tree x d)
  117.   (rotate+! tree x (-d d)))
  118.  
  119. (define (rb-tree/insert! tree key datum)
  120.   (guarantee-rb-tree tree 'RB-TREE/INSERT!)
  121.   (let ((key=? (tree-key=? tree))
  122.     (key<? (tree-key<? tree)))
  123.     (let loop ((x (tree-root tree)) (y #f) (d #f))
  124.       (cond ((not x)
  125.          (let ((z (make-node key datum)))
  126.            (without-interrupts
  127.         (lambda ()
  128.           (set-node-up! z y)
  129.           (cond ((not y) (set-tree-root! tree z))
  130.             ((eq? 'LEFT d) (set-node-left! y z))
  131.             (else (set-node-right! y z)))
  132.           (set-node-color! z 'RED)
  133.           (insert-fixup! tree z)))))
  134.         ((key=? key (node-key x)) (set-node-datum! x datum))
  135.         ((key<? key (node-key x)) (loop (node-left x) x 'LEFT))
  136.         (else (loop (node-right x) x 'RIGHT))))))
  137.  
  138. (define (insert-fixup! tree x)
  139.   ;; Assumptions: X is red, and the only possible violation of the
  140.   ;; tree properties is that (NODE-UP X) is also red.
  141.   (let loop ((x x))
  142.     (let ((u (node-up x)))
  143.       (if (and u (eq? 'RED (node-color u)))
  144.       (let ((d (b->d (eq? u (node-left (node-up u))))))
  145.         (let ((y (get-link- (node-up u) d)))
  146.           (if (and y (eq? 'RED (node-color y)))
  147.           ;; case 1
  148.           (begin
  149.             (set-node-color! u 'BLACK)
  150.             (set-node-color! y 'BLACK)
  151.             (set-node-color! (node-up u) 'RED)
  152.             (loop (node-up u)))
  153.           (let ((x
  154.              (if (eq? x (get-link- u d))
  155.                  ;; case 2
  156.                  (begin
  157.                    (rotate+! tree u d)
  158.                    u)
  159.                  x)))
  160.             ;; case 3
  161.             (let ((u (node-up x)))
  162.               (set-node-color! u 'BLACK)
  163.               (set-node-color! (node-up u) 'RED)
  164.               (rotate-! tree (node-up u) d)))))))))
  165.   (set-node-color! (tree-root tree) 'BLACK))
  166.  
  167. (define (alist->rb-tree alist key=? key<?)
  168.   ;; Is there a more efficient way to do this?
  169.   (let ((tree (make-rb-tree key=? key<?)))
  170.     (do ((alist alist (cdr alist)))
  171.     ((null? alist))
  172.       (rb-tree/insert! tree (caar alist) (cdar alist)))
  173.     tree))
  174.  
  175. (define-integrable (without-interrupts thunk)
  176.   (let ((interrupt-mask (set-interrupt-enables! interrupt-mask/gc-ok)))
  177.     (thunk)
  178.     (set-interrupt-enables! interrupt-mask)
  179.     unspecific))
  180.  
  181. (define (rb-tree/delete! tree key)
  182.   (guarantee-rb-tree tree 'RB-TREE/DELETE!)
  183.   (let ((key=? (tree-key=? tree))
  184.     (key<? (tree-key<? tree)))
  185.     (let loop ((x (tree-root tree)))
  186.       (cond ((not x) unspecific)
  187.         ((key=? key (node-key x)) (delete-node! tree x))
  188.         ((key<? key (node-key x)) (loop (node-left x)))
  189.         (else (loop (node-right x)))))))
  190.  
  191. (define (delete-node! tree z)
  192.   (without-interrupts
  193.    (lambda ()
  194.      (let ((z
  195.         (if (and (node-left z) (node-right z))
  196.         (let ((y (next-node z)))
  197.           (set-node-key! z (node-key y))
  198.           (set-node-datum! z (node-datum y))
  199.           y)
  200.         z)))
  201.        (let ((x (or (node-left z) (node-right z)))
  202.          (u (node-up z)))
  203.      (if x (set-node-up! x u))
  204.      (cond ((not u) (set-tree-root! tree x))
  205.            ((eq? z (node-left u)) (set-node-left! u x))
  206.            (else (set-node-right! u x)))
  207.      (if (eq? 'BLACK (node-color z))
  208.          (delete-fixup! tree x u)))))))
  209.  
  210. (define (delete-fixup! tree x u)
  211.   (let loop ((x x) (u u))
  212.     (if (or (not u)
  213.         (and x (eq? 'RED (node-color x))))
  214.     (if x (set-node-color! x 'BLACK))
  215.     (let ((d (b->d (eq? x (node-left u)))))
  216.       (let ((w
  217.          (let ((w (get-link- u d)))
  218.            (if (eq? 'RED (node-color w))
  219.                ;; case 1
  220.                (begin
  221.              (set-node-color! w 'BLACK)
  222.              (set-node-color! u 'RED)
  223.              (rotate+! tree u d)
  224.              (get-link- u d))
  225.                w)))
  226.         (case-4
  227.          (lambda (w)
  228.            (set-node-color! w (node-color u))
  229.            (set-node-color! u 'BLACK)
  230.            (set-node-color! (get-link- w d) 'BLACK)
  231.            (rotate+! tree u d)
  232.            (set-node-color! (tree-root tree) 'BLACK))))
  233.         (if (let ((n- (get-link- w d)))
  234.           (and n-
  235.                (eq? 'RED (node-color n-))))
  236.         (case-4 w)
  237.         (let ((n+ (get-link+ w d)))
  238.           (if (or (not n+)
  239.               (eq? 'BLACK (node-color n+)))
  240.               ;; case 2
  241.               (begin
  242.             (set-node-color! w 'RED)
  243.             (loop u (node-up u)))
  244.               ;; case 3
  245.               (begin
  246.             (set-node-color! n+ 'BLACK)
  247.             (set-node-color! w 'RED)
  248.             (rotate-! tree w d)
  249.             (case-4 (get-link- u d)))))))))))
  250.  
  251. (define (rb-tree/lookup tree key default)
  252.   (guarantee-rb-tree tree 'RB-TREE/LOOKUP)
  253.   (let ((key=? (tree-key=? tree))
  254.     (key<? (tree-key<? tree)))
  255.     (let loop ((x (tree-root tree)))
  256.       (cond ((not x) default)
  257.         ((key=? key (node-key x)) (node-datum x))
  258.         ((key<? key (node-key x)) (loop (node-left x)))
  259.         (else (loop (node-right x)))))))
  260.  
  261. (define (rb-tree/copy tree)
  262.   (guarantee-rb-tree tree 'RB-TREE/COPY)
  263.   (let ((result (make-rb-tree (tree-key=? tree) (tree-key<? tree))))
  264.     (set-tree-root!
  265.      result
  266.      (let loop ((node (tree-root tree)) (up #f))
  267.        (and node
  268.         (let ((node* (make-node (node-key node) (node-datum node))))
  269.           (set-node-color! node* (node-color node))
  270.           (set-node-up! node* up)
  271.           (set-node-left! node* (loop (node-left node) node*))
  272.           (set-node-right! node* (loop (node-right node) node*))
  273.           node*))))
  274.     result))
  275.  
  276. (define (rb-tree/height tree)
  277.   (guarantee-rb-tree tree 'RB-TREE/HEIGHT)
  278.   (let loop ((node (tree-root tree)))
  279.     (if node
  280.     (+ 1 (max (loop (node-left node)) (loop (node-right node))))
  281.     0)))
  282.  
  283. (define (rb-tree/size tree)
  284.   (guarantee-rb-tree tree 'RB-TREE/SIZE)
  285.   (let loop ((node (tree-root tree)))
  286.     (if node
  287.     (+ 1 (loop (node-left node)) (loop (node-right node)))
  288.     0)))
  289.  
  290. (define (rb-tree/empty? tree)
  291.   (guarantee-rb-tree tree 'RB-TREE/EMPTY?)
  292.   (not (tree-root tree)))
  293.  
  294. (define (rb-tree/equal? x y datum=?)
  295.   (guarantee-rb-tree x 'RB-TREE/EQUAL?)
  296.   (guarantee-rb-tree y 'RB-TREE/EQUAL?)
  297.   (let ((key=? (tree-key=? x)))
  298.     (and (eq? key=? (tree-key=? y))
  299.      (let loop ((nx (min-node x)) (ny (min-node y)))
  300.        (if (not nx)
  301.            (not ny)
  302.            (and ny
  303.             (key=? (node-key nx) (node-key ny))
  304.             (datum=? (node-datum nx) (node-datum ny))
  305.             (loop (next-node nx) (next-node ny))))))))
  306.  
  307. (define (rb-tree->alist tree)
  308.   (guarantee-rb-tree tree 'RB-TREE->ALIST)
  309.   (let ((node (min-node tree)))
  310.     (if node
  311.     (let ((result (list (cons (node-key node) (node-datum node)))))
  312.       (let loop ((node (next-node node)) (prev result))
  313.         (if node
  314.         (let ((pair (list (cons (node-key node) (node-datum node)))))
  315.           (set-cdr! prev pair)
  316.           (loop (next-node node) pair))))
  317.       result)
  318.     '())))
  319.  
  320. (define (rb-tree/key-list tree)
  321.   (guarantee-rb-tree tree 'RB-TREE/KEY-LIST)
  322.   (let ((node (min-node tree)))
  323.     (if node
  324.     (let ((result (list (node-key node))))
  325.       (let loop ((node (next-node node)) (prev result))
  326.         (if node
  327.         (let ((pair (list (node-key node))))
  328.           (set-cdr! prev pair)
  329.           (loop (next-node node) pair))))
  330.       result)
  331.     '())))
  332.  
  333. (define (rb-tree/datum-list tree)
  334.   (guarantee-rb-tree tree 'RB-TREE/DATUM-LIST)
  335.   (let ((node (min-node tree)))
  336.     (if node
  337.     (let ((result (list (node-datum node))))
  338.       (let loop ((node (next-node node)) (prev result))
  339.         (if node
  340.         (let ((pair (list (node-datum node))))
  341.           (set-cdr! prev pair)
  342.           (loop (next-node node) pair))))
  343.       result)
  344.     '())))
  345.  
  346. (define (rb-tree/min tree default)
  347.   (guarantee-rb-tree tree 'RB-TREE/MIN)
  348.   (let ((node (min-node tree)))
  349.     (if node
  350.     (node-key node)
  351.     default)))
  352.  
  353. (define (rb-tree/min-datum tree default)
  354.   (guarantee-rb-tree tree 'RB-TREE/MIN-DATUM)
  355.   (let ((node (min-node tree)))
  356.     (if node
  357.     (node-datum node)
  358.     default)))
  359.  
  360. (define (rb-tree/min-pair tree)
  361.   (guarantee-rb-tree tree 'RB-TREE/MIN-PAIR)
  362.   (let ((node (min-node tree)))
  363.     (and node
  364.      (node-pair node))))
  365.  
  366. (define (rb-tree/delete-min! tree default)
  367.   (guarantee-rb-tree tree 'RB-TREE/DELETE-MIN!)
  368.   (let ((node (min-node tree)))
  369.     (if node
  370.     (let ((key (node-key node)))
  371.       (delete-node! tree node)
  372.       key)
  373.     default)))
  374.  
  375. (define (rb-tree/delete-min-datum! tree default)
  376.   (guarantee-rb-tree tree 'RB-TREE/DELETE-MIN-DATUM!)
  377.   (let ((node (min-node tree)))
  378.     (if node
  379.     (let ((datum (node-datum node)))
  380.       (delete-node! tree node)
  381.       datum)
  382.     default)))
  383.  
  384. (define (rb-tree/delete-min-pair! tree)
  385.   (guarantee-rb-tree tree 'RB-TREE/DELETE-MIN-PAIR!)
  386.   (let ((node (min-node tree)))
  387.     (and node
  388.      (let ((pair (node-pair node)))
  389.        (delete-node! tree node)
  390.        pair))))
  391.  
  392. (define (rb-tree/max tree default)
  393.   (guarantee-rb-tree tree 'RB-TREE/MAX)
  394.   (let ((node (max-node tree)))
  395.     (if node
  396.     (node-key node)
  397.     default)))
  398.  
  399. (define (rb-tree/max-datum tree default)
  400.   (guarantee-rb-tree tree 'RB-TREE/MAX-DATUM)
  401.   (let ((node (max-node tree)))
  402.     (if node
  403.     (node-datum node)
  404.     default)))
  405.  
  406. (define (rb-tree/max-pair tree)
  407.   (guarantee-rb-tree tree 'RB-TREE/MAX-PAIR)
  408.   (let ((node (max-node tree)))
  409.     (and node
  410.      (node-pair node))))
  411.  
  412. (define (rb-tree/delete-max! tree default)
  413.   (guarantee-rb-tree tree 'RB-TREE/DELETE-MAX!)
  414.   (let ((node (max-node tree)))
  415.     (if node
  416.     (let ((key (node-key node)))
  417.       (delete-node! tree node)
  418.       key)
  419.     default)))
  420.  
  421. (define (rb-tree/delete-max-datum! tree default)
  422.   (guarantee-rb-tree tree 'RB-TREE/DELETE-MAX-DATUM!)
  423.   (let ((node (max-node tree)))
  424.     (if node
  425.     (let ((datum (node-datum node)))
  426.       (delete-node! tree node)
  427.       datum)
  428.     default)))
  429.  
  430. (define (rb-tree/delete-max-pair! tree)
  431.   (guarantee-rb-tree tree 'RB-TREE/DELETE-MAX-PAIR!)
  432.   (let ((node (max-node tree)))
  433.     (and node
  434.      (let ((pair (node-pair node)))
  435.        (delete-node! tree node)
  436.        pair))))
  437.  
  438. (define (min-node tree)
  439.   (and (tree-root tree)
  440.        (let loop ((x (tree-root tree)))
  441.      (if (node-left x)
  442.          (loop (node-left x))
  443.          x))))
  444.  
  445. (define (max-node tree)
  446.   (and (tree-root tree)
  447.        (let loop ((x (tree-root tree)))
  448.      (if (node-right x)
  449.          (loop (node-right x))
  450.          x))))
  451.  
  452. (define (next-node x)
  453.   (if (node-right x)
  454.        (let loop ((x (node-right x)))
  455.      (if (node-left x)
  456.          (loop (node-left x))
  457.          x))
  458.       (let loop ((x x))
  459.     (let ((y (node-up x)))
  460.       (if (and y (eq? x (node-right y)))
  461.           (loop y)
  462.           y)))))
  463.  
  464. (define-integrable (node-pair node)
  465.   (cons (node-key node) (node-datum node)))