442 lines
14 KiB
Scheme
442 lines
14 KiB
Scheme
; Copyright (c) 1993-1999 by Richard Kelsey and Jonathan Rees. See file COPYING.
|
|
|
|
; Red-Black binary search trees as described in Introduction to Algorithms
|
|
; by Cormen, Leiserson, and Rivest. Look there if you want to understand
|
|
; the algorithm.
|
|
;
|
|
; These are like tables in that the value of a key defaults to #f.
|
|
;
|
|
; (make-search-tree key-= key-<) -> tree
|
|
;
|
|
; (search-tree? value) -> boolean
|
|
;
|
|
; (search-tree-ref tree key) -> value
|
|
;
|
|
; (search-tree-set! tree key value)
|
|
;
|
|
; (search-tree-modify! tree key proc)
|
|
; == (search-tree-set! tree key (proc (search-tree-ref tree key)))
|
|
;
|
|
; (search-tree-max tree) -> key + value
|
|
; (pop-search-tree-max! tree) -> key + value (removes entry)
|
|
;
|
|
; (search-tree-min tree) -> key + value
|
|
; (pop-search-tree-min! tree) -> key + value (removes entry)
|
|
;
|
|
; (walk-search-tree proc tree)
|
|
; applies PROC in order to all key + value pairs with a non-#f value
|
|
|
|
(define-record-type tree :tree
|
|
(make-tree lookup nil root)
|
|
search-tree?
|
|
(lookup tree-lookup)
|
|
(nil tree-nil) ; node marker for missing leaf nodes
|
|
(root tree-root set-tree-root!))
|
|
|
|
(define (make-search-tree = <)
|
|
(let ((nil (make-node #f #f #f)))
|
|
(set-node-red?! nil #f)
|
|
(make-tree (make-lookup = <) nil #f)))
|
|
|
|
(define-record-type node :node
|
|
(really-make-node key value parent red? left right)
|
|
node?
|
|
(key node-key set-node-key!)
|
|
(value node-value set-node-value!)
|
|
(parent node-parent set-node-parent!) ; #f in the root node
|
|
(red? node-red? set-node-red?!) ; for balancing the tree
|
|
(left node-left set-node-left!) ; left and
|
|
(right node-right set-node-right!)) ; right subtrees
|
|
|
|
(define (make-node key value parent)
|
|
(really-make-node key value parent #t #f #f))
|
|
|
|
(define-record-discloser :node
|
|
(lambda (node)
|
|
(list 'node (node-key node))))
|
|
|
|
; Lookup up KEY and return its value.
|
|
|
|
(define (search-tree-ref tree key)
|
|
(call-with-values
|
|
(lambda ()
|
|
((tree-lookup tree) tree key))
|
|
(lambda (node parent left?)
|
|
(if node
|
|
(node-value node)
|
|
#f))))
|
|
|
|
; Adding and modifying entries.
|
|
|
|
(define (search-tree-set! tree key value)
|
|
(search-tree-modify! tree key (lambda (ignore) value)))
|
|
|
|
(define (search-tree-modify! tree key proc)
|
|
(call-with-values
|
|
(lambda ()
|
|
((tree-lookup tree) tree key))
|
|
(lambda (node parent left?)
|
|
(let ((new-value (proc (if node (node-value node) #f))))
|
|
(cond ((and node new-value)
|
|
(set-node-value! node new-value))
|
|
(new-value
|
|
(really-insert! tree parent left? (make-node key new-value parent)))
|
|
(node
|
|
(really-delete! tree node)))))))
|
|
|
|
; Min and max entries.
|
|
|
|
(define (search-tree-max tree)
|
|
(real-search-tree-max tree #f))
|
|
|
|
(define (pop-search-tree-max! tree)
|
|
(real-search-tree-max tree #t))
|
|
|
|
(define (real-search-tree-max tree delete?)
|
|
(let ((node (tree-root tree)))
|
|
(if node
|
|
(let loop ((node node))
|
|
(cond ((node-right node)
|
|
=> loop)
|
|
(else
|
|
(if delete?
|
|
(really-delete! tree node))
|
|
(values (node-key node) (node-value node)))))
|
|
(values #f #f))))
|
|
|
|
(define (search-tree-min tree)
|
|
(real-search-tree-min tree #f))
|
|
|
|
(define (pop-search-tree-min! tree)
|
|
(real-search-tree-min tree #t))
|
|
|
|
(define (real-search-tree-min tree delete?)
|
|
(let ((node (tree-root tree)))
|
|
(if node
|
|
(let loop ((node node))
|
|
(cond ((node-left node)
|
|
=> loop)
|
|
(else
|
|
(if delete?
|
|
(really-delete! tree node))
|
|
(values (node-key node) (node-value node)))))
|
|
(values #f #f))))
|
|
|
|
(define (walk-search-tree proc tree)
|
|
(let recur ((node (tree-root tree)))
|
|
(cond (node
|
|
(recur (node-left node))
|
|
(proc (node-key node) (node-value node))
|
|
(recur (node-right node))))))
|
|
|
|
; Lookup up an entry. Easy.
|
|
;
|
|
; Hack of checking common case reduced lookup time in a 1000 element search
|
|
; tree by a third.
|
|
|
|
(define (make-lookup tree-= tree-<)
|
|
(if (and (eq? tree-= =)
|
|
(eq? tree-< <))
|
|
default-lookup
|
|
(lambda (tree key)
|
|
(let loop ((node (tree-root tree))
|
|
(parent #f)
|
|
(left? #f))
|
|
(cond ((not node)
|
|
(values #f parent left?))
|
|
((tree-= (node-key node) key)
|
|
(values node #f #f))
|
|
((tree-< key (node-key node))
|
|
(loop (node-left node) node #t))
|
|
(else
|
|
(loop (node-right node) node #f)))))))
|
|
|
|
(define (default-lookup tree key)
|
|
(let loop ((node (tree-root tree))
|
|
(parent #f)
|
|
(left? #f))
|
|
(cond ((not node)
|
|
(values #f parent left?))
|
|
((= (node-key node) key)
|
|
(values node #f #f))
|
|
((< key (node-key node))
|
|
(loop (node-left node) node #t))
|
|
(else
|
|
(loop (node-right node) node #f)))))
|
|
|
|
;----------------------------------------------------------------
|
|
; Little utilities.
|
|
|
|
; Parameterized node access
|
|
|
|
(define (node-child node left?)
|
|
(if left?
|
|
(node-left node)
|
|
(node-right node)))
|
|
|
|
(define (set-node-child! node left? child)
|
|
(if left?
|
|
(set-node-left! node child)
|
|
(set-node-right! node child)))
|
|
|
|
; Empty leaf slots are considered black.
|
|
|
|
(define (node-black? node)
|
|
(not (and node (node-red? node))))
|
|
|
|
; The next node (used in REALLY-DELETE!)
|
|
|
|
(define (successor node)
|
|
(cond ((node-right node)
|
|
=> (lambda (node)
|
|
(let loop ((node node))
|
|
(cond ((node-left node)
|
|
=> loop)
|
|
(else node)))))
|
|
(else
|
|
(let loop ((node node) (parent (node-parent node)))
|
|
(if (and parent
|
|
(eq? node (node-right parent)))
|
|
(loop parent (node-parent parent))
|
|
parent)))))
|
|
|
|
;----------------------------------------------------------------
|
|
; Add NODE as the LEFT? child of PARENT and balance the tree.
|
|
|
|
(define (really-insert! tree parent left? node)
|
|
(if (not parent)
|
|
(set-tree-root! tree node)
|
|
(set-node-child! parent left? node))
|
|
(fixup-insertion! node tree))
|
|
|
|
; Balance the tree after NODE has been inserted.
|
|
|
|
(define (fixup-insertion! node tree)
|
|
(let loop ((node node))
|
|
(let ((parent (node-parent node)))
|
|
(if (and parent (node-red? parent))
|
|
(let* ((grand (node-parent parent))
|
|
(left? (eq? parent (node-left grand)))
|
|
(y (node-child grand (not left?))))
|
|
(cond ((node-black? y)
|
|
(let* ((node (cond ((eq? node (node-child parent (not left?)))
|
|
(rotate! parent left? tree)
|
|
parent)
|
|
(else node)))
|
|
(parent (node-parent node))
|
|
(grand (node-parent parent)))
|
|
(set-node-red?! parent #f)
|
|
(set-node-red?! grand #t)
|
|
(rotate! grand (not left?) tree)
|
|
(loop node)))
|
|
(else
|
|
(set-node-red?! parent #f)
|
|
(set-node-red?! y #f)
|
|
(set-node-red?! grand #t)
|
|
(loop grand)))))))
|
|
(set-node-red?! (tree-root tree) #f))
|
|
|
|
; A B
|
|
; / \ =(rotate! A #f tree)=> / \
|
|
; B k i A
|
|
; / \ <=(rotate! B #t tree)= / \
|
|
; i j j k
|
|
|
|
(define (rotate! node left? tree)
|
|
(let* ((y (node-child node (not left?)))
|
|
(y-left (node-child y left?))
|
|
(parent (node-parent node)))
|
|
(set-node-child! node (not left?) y-left)
|
|
(if y-left
|
|
(set-node-parent! y-left node))
|
|
(replace! parent y node tree)
|
|
(set-node-child! y left? node)
|
|
(set-node-parent! node y)))
|
|
|
|
; Replace CHILD (of PARENT) with NEW-CHILD
|
|
|
|
(define (replace! parent new-child child tree)
|
|
(set-node-parent! new-child parent)
|
|
(cond ((eq? child (tree-root tree))
|
|
(set-tree-root! tree new-child))
|
|
((eq? child (node-left parent))
|
|
(set-node-left! parent new-child))
|
|
(else
|
|
(set-node-right! parent new-child))))
|
|
|
|
; Remove NODE from tree.
|
|
|
|
(define (really-delete! tree node)
|
|
(let* ((y (cond ((or (not (node-left node))
|
|
(not (node-right node)))
|
|
node)
|
|
(else
|
|
(let ((y (successor node)))
|
|
(set-node-key! node (node-key y))
|
|
(set-node-value! node (node-value y))
|
|
y))))
|
|
(x (or (node-left y)
|
|
(node-right y)
|
|
(let ((x (tree-nil tree)))
|
|
(set-node-right! y x)
|
|
x)))
|
|
(parent (node-parent y)))
|
|
(replace! parent x y tree)
|
|
(if (not (node-red? y))
|
|
(fixup-delete! x tree))
|
|
(let ((nil (tree-nil tree)))
|
|
(cond ((node-parent nil)
|
|
=> (lambda (p)
|
|
(if (eq? (node-right p) nil)
|
|
(set-node-right! p #f)
|
|
(set-node-left! p #f))
|
|
(set-node-parent! (tree-nil tree) #f)))
|
|
((eq? nil (tree-root tree))
|
|
(set-tree-root! tree #f))))))
|
|
|
|
(define (fixup-delete! x tree)
|
|
(let loop ((x x))
|
|
(if (or (eq? x (tree-root tree))
|
|
(node-red? x))
|
|
(set-node-red?! x #f)
|
|
(let* ((parent (node-parent x))
|
|
(left? (eq? x (node-left parent)))
|
|
(w (node-child parent (not left?)))
|
|
(w (cond ((node-red? w)
|
|
(set-node-red?! w #f)
|
|
(set-node-red?! parent #t)
|
|
(rotate! parent left? tree)
|
|
(node-child (node-parent x) (not left?)))
|
|
(else
|
|
w))))
|
|
(cond ((and (node-black? (node-left w))
|
|
(node-black? (node-right w)))
|
|
(set-node-red?! w #t)
|
|
(loop (node-parent x)))
|
|
(else
|
|
(let ((w (cond ((node-black? (node-child w (not left?)))
|
|
(set-node-red?! (node-child w left?) #f)
|
|
(set-node-red?! w #t)
|
|
(rotate! w (not left?) tree)
|
|
(node-child (node-parent x) (not left?)))
|
|
(else
|
|
w))))
|
|
(let ((parent (node-parent x)))
|
|
(set-node-red?! w (node-red? parent))
|
|
(set-node-red?! parent #f)
|
|
(set-node-red?! (node-child w (not left?)) #f)
|
|
(rotate! parent left? tree)
|
|
(set-node-red?! (tree-root tree) #f)))))))))
|
|
|
|
; Verify that the coloring is correct
|
|
;
|
|
;(define (okay-tree? tree)
|
|
; (receive (okay? red? count)
|
|
; (let recur ((node (tree-root tree)))
|
|
; (if (not node)
|
|
; (values #t #f 0)
|
|
; (receive (l-ok? l-r? l-c)
|
|
; (recur (node-left node))
|
|
; (receive (r-ok? r-r? r-c)
|
|
; (recur (node-right node))
|
|
; (values (and l-ok?
|
|
; r-ok?
|
|
; (not (and (node-red? node)
|
|
; (or l-r? r-r?)))
|
|
; (= l-c r-c))
|
|
; (node-red? node)
|
|
; (if (node-red? node)
|
|
; l-c
|
|
; (+ l-c 1)))))))
|
|
; okay?))
|
|
;
|
|
;
|
|
;(define (walk-sequences proc list)
|
|
; (let recur ((list list) (r '()))
|
|
; (if (null? list)
|
|
; (proc (reverse r))
|
|
; (let loop ((list list) (done '()))
|
|
; (if (not (null? list))
|
|
; (let ((next (car list)))
|
|
; (recur (append (reverse done) (cdr list)) (cons next r))
|
|
; (loop (cdr list) (cons next done))))))))
|
|
;
|
|
;(define (tree-test n)
|
|
; (let ((iota (do ((i n (- i 1))
|
|
; (l '() (cons i l)))
|
|
; ((<= i 0) l))))
|
|
; (walk-sequences (lambda (in)
|
|
; (walk-sequences (lambda (out)
|
|
; (do-tree-test in out))
|
|
; iota))
|
|
; iota)
|
|
; #t))
|
|
;
|
|
;(define (do-tree-test in out)
|
|
; (let ((tree (make-search-tree = <)))
|
|
; (for-each (lambda (i)
|
|
; (search-tree-set! tree i (- 0 i)))
|
|
; in)
|
|
; (if (not (okay-tree? tree))
|
|
; (breakpoint "tree ~S is not okay" in))
|
|
; (if (not (tree-ordered? tree (length in)))
|
|
; (breakpoint "tree ~S is not ordered" in))
|
|
; (for-each (lambda (i)
|
|
; (if (not (= (search-tree-ref tree i) (- 0 i)))
|
|
; (breakpoint "looking up ~S in ~S lost" i in)))
|
|
; in)
|
|
; (do ((o out (cdr o)))
|
|
; ((null? o))
|
|
; (search-tree-set! tree (car o) #f)
|
|
; (if (not (okay-tree? tree))
|
|
; (breakpoint "tree ~S is not okay after deletions ~S" in out)))))
|
|
;
|
|
;(define (tree-ordered? tree count)
|
|
; (let ((l '()))
|
|
; (walk-search-tree (lambda (key value)
|
|
; (set! l (cons (cons key value) l)))
|
|
; tree)
|
|
; (let loop ((l l) (n count))
|
|
; (cond ((null? l)
|
|
; (= n 0))
|
|
; ((and (= (caar l) n)
|
|
; (= (cdar l) (- 0 n)))
|
|
; (loop (cdr l) (- n 1)))
|
|
; (else #f)))))
|
|
;
|
|
;(define (do-tests tester)
|
|
; (do ((i 0 (+ i 1)))
|
|
; (#f)
|
|
; (tester i)
|
|
; (format #t " done with ~D~%" i)))
|
|
;
|
|
;(define (another-test n)
|
|
; (let ((iota (do ((i n (- i 1))
|
|
; (l '() (cons i l)))
|
|
; ((<= i 0) l))))
|
|
; (walk-sequences (lambda (in)
|
|
; (do ((i 1 (+ i 1)))
|
|
; ((> i n))
|
|
; (let ((tree (make-search-tree = <)))
|
|
; (for-each (lambda (i)
|
|
; (search-tree-set! tree i (- 0 i)))
|
|
; in)
|
|
; (if (not (okay-tree? tree))
|
|
; (breakpoint "tree ~S is not okay" in))
|
|
; (if (not (tree-ordered? tree (length in)))
|
|
; (breakpoint "tree ~S is not ordered" in))
|
|
; (for-each (lambda (i)
|
|
; (if (not (= (search-tree-ref tree i) (- 0 i)))
|
|
; (breakpoint "looking up ~S in ~S lost" i in)))
|
|
; in)
|
|
; (search-tree-set! tree i #f)
|
|
; (if (not (okay-tree? tree))
|
|
; (breakpoint "tree ~S is not okay after deletion ~S"
|
|
; in i))
|
|
; (for-each (lambda (j)
|
|
; (let ((ref (search-tree-ref tree j)))
|
|
; (if (not (eq? ref (if (= j i) #f (- 0 j))))
|
|
; (breakpoint "looking up ~S in ~S lost" i in))))
|
|
; in))))
|
|
; iota)))
|