;; ;; Red-black tree ;; ;; An implementation of an ordered dictionary data structure, based ;; on red-black trees. ;; ;; This code is based on the SML/NJ library implementation of ;; red-black trees, which is in turn based on Chris Okasaki's ;; implementation of red-black trees. The delete function is based on ;; the description in Cormen, Leiserson, and Rivest. ;; ;; Some helper code was borrowed from treap.scm by Oleg Kiselyov. ;; ;; ;; Copyright 2007-2018 Ivan Raikov. ;; ;; ;; This program is free software: you can redistribute it and/or ;; modify it under the terms of the GNU General Public License as ;; published by the Free Software Foundation, either version 3 of the ;; License, or (at your option) any later version. ;; ;; This program is distributed in the hope that it will be useful, but ;; WITHOUT ANY WARRANTY; without even the implied warranty of ;; MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU ;; General Public License for more details. ;; ;; A full copy of the GPL license can be found at ;; . ;; (module rb-tree ( rb-tree-map union-with union-withi persistent-map? get get/default get-min get-max get-value get-value/default put generate update delete for-each-ascending for-each-descending) (import scheme (chicken base) (chicken memory representation) (only (chicken string) conc) datatype matchable yasos yasos-collections) (import-for-syntax srfi-1) ;; ;; Persistent map operations ;; ;; A persistent map class provides the operations provided by yasos-collections, ;; plus the following operations: ;; ;; get ;; ;; procedure LAMBDA KEY . DEFAULT-CLAUSE which searches ;; the given tree for an association with a given KEY, and ;; returns a (key . value) pair of the found association. ;; If an association with the KEY cannot be located in ;; the tree, the PROC returns the result of evaluating the ;; DEFAULT-CLAUSE. If the default clause is omitted, an ;; error is signalled. The KEY must be comparable to the ;; keys in the map by a key-compare predicate (which has ;; been specified when the map was created) ;; ;; get-min ;; ;; returns a (key . value) pair for an association in the ;; tree with the smallest key. If the tree is empty, ;; an error is signalled. ;; ;; get-max ;; ;; returns a (key . value) pair for an association in the ;; tree with the largest key. If the tree is empty, an ;; error is signalled. ;; ;; ;; for-each-ascending ;; ;; a procedure LAMBDA PROC that will apply the given ;; procedure PROC to each (key . value) association of ;; the tree, from the one with the smallest key all the ;; way to the one with the max key, in an ascending order ;; of keys. The tree must not be empty. ;; ;; for-each-descending ;; ;; a procedure LAMBDA PROC that will apply the given ;; procedure PROC to each (key . value) association of ;; the tree, in the descending order of keys. ;; The tree must not be empty. ;; (define-predicate persistent-map?) (define-operation (get pmap key)) (define-operation (get/default pmap key default-clause)) (define-operation (get-min pmap)) (define-operation (get-max pmap)) (define-operation (get-value pmap key)) (define-operation (get-value/default pmap key default-clause)) (define-operation (generate pmap p f g seed)) (define-operation (put pmap key value)) (define-operation (update pmap key value merge-fn)) (define-operation (delete pmap key)) (define-operation (for-each-ascending pmap f)) (define-operation (for-each-descending pmap f)) ;; ;; A red-black tree should satisfy the following two invariants: ;; ;; Red Invariant: each red node has a black parent. ;; ;; Black Condition: each path from the root to an empty node has the ;; same number of black nodes (the tree's black height). ;; ;; The Red condition implies that the root is always black and the Black ;; condition implies that any node with only one child will be black and ;; its child will be a red leaf. ;; ;; ;; The red-black tree object is created by procedure ;; rb-tree-map ;; ;; rb-tree-tree-persistent-map :: KEY-COMPARE-PROC -> PERSISTENT-MAP ;; ;; where KEY-COMPARE-PROC is a user-supplied function ;; ;; KEY-COMPARE-PROC:: key1 key2 -> INTEGER ;; ;; that takes two keys and returns a negative, positive, or zero ;; number depending on how the first key compares to the second. ;; (define (rb-tree:error x . rest) (let ((port (open-output-string))) (let loop ((objs (cons x rest))) (if (null? objs) (begin (newline port) (error 'rb-tree (get-output-string port))) (begin (display (car objs) port) (display " " port) (loop (cdr objs))))))) (define R 'Red) (define B 'Black) (define (color? x) (or (eq? x 'Red) (eq? x 'Black))) (define-datatype tree tree? (Empty) (Tree (color color?) (left tree?) (key identity) (value identity) (right tree?))) (define-datatype zipper zipper? (Top) (Left (color color?) (key identity) (value identity) (tree tree?) (zipper zipper?)) (Right (color color?) (tree tree?) (key identity) (value identity) (zipper zipper?))) (define (tree-tag x) (cases tree x (Empty () 'Empty) (Tree (c l k v r) 'Tree))) (define-record-printer (tree x out) (cases tree x (Empty () (display "#(Empty)" out)) (Tree (c l k v r) (display "#(Tree " out) (display (conc c " ") out) (display (tree-tag l) out) (display (conc " " k ":" v " ") out) (display (tree-tag r) out) (display ")" out)))) ;; ;; This macro was borrowed from treap.scm by Oleg Kiselyov ;; (define-syntax dispatch-on-key (er-macro-transformer (lambda (x r c) (let ((key-compare (second x)) (key (third x)) (node-key (fourth x)) (on-less (fifth x)) (on-equal (sixth x)) (on-greater (seventh x))) (let ((%let (r 'let)) (%cond (r 'cond)) (%else (r 'else)) (%zero? (r 'zero?)) (%positive? (r 'positive?)) (result (r 'result))) `(,%let ((,result (,key-compare ,key ,node-key ))) (,%cond ((,%zero? ,result) ,on-equal) ((,%positive? ,result) ,on-greater) (,%else ,on-less)))))))) ;; Adds a new association to the tree (or replaces the old one if existed). (define (insert key-compare value-merge root key value) (let ins ((root root)) (cases tree root (Empty () (Tree R (Empty) key value (Empty))) (Tree (color a yk y b) (dispatch-on-key key-compare key yk ;; Case 1: key < yk (match a (($ rb-tree#tree 'Tree 'Red c zk z d) (dispatch-on-key key-compare key zk ;; Case 1.1: key < zk (let ((c1 (ins c))) (match c1 (($ rb-tree#tree 'Tree 'Red e wk w f) (Tree R (Tree B e wk w f) zk z (Tree B d yk y b))) (else (Tree B (Tree R c1 zk z d) yk y b)))) ;; Case 1.2: key = zk (Tree color (Tree R c key (value-merge value z) d) yk y b) ;; Case 1.3: key > zk (let ((d1 (ins d))) (match d1 (($ rb-tree#tree 'Tree 'Red e wk w f) (Tree R (Tree B c zk z e) wk w (Tree B f yk y b))) (else (Tree B (Tree R c zk z d1) yk y b)))))) (else (let ((a1 (ins a))) (Tree B a1 yk y b)))) ;; Case 2: key = yk (Tree color a key (value-merge value y) b) ;; Case 3: key > yk (match b (($ rb-tree#tree 'Tree 'Red c zk z d) (dispatch-on-key key-compare key zk ;; Case 3.1: key < zk (let ((c1 (ins c))) (match c1 (($ rb-tree#tree 'Tree 'Red e wk w f) (Tree R (Tree B a yk y e) wk w (Tree B f zk z d))) (else (Tree B a yk y (Tree R c1 zk z d))))) ;; Case 3.2: key = zk (Tree color a yk y (Tree R c key (value-merge value z) d)) ;; Case 3.3: key > zk (let ((d1 (ins d))) (match d1 (($ rb-tree#tree 'Tree 'Red e wk w f) (Tree R (Tree B a yk y c) zk z (Tree B e wk w f))) (else (Tree B a yk y (Tree R c zk z d1))))))) (else (let ((b1 (ins b))) (Tree B a yk y b1)))))))) ) ;; Looks for an item: Given a key, returns the corresponding (key ;; . value) association or #f if the tree does not contain an ;; association with that key. (define (find-assoc key-compare root key) (let recur ((root root)) (cases tree root (Empty () #f) (Tree (c a yk y b) (dispatch-on-key key-compare key yk (recur a) (cons yk y) (recur b))))) ) ;; Looks for an item: Given a key, returns the value of the ;; corresponding (key . value) association or #f if the tree does ;; not contain an association with that key. (define (find-ref key-compare root key) (let recur ((root root)) (cases tree root (Empty () #f) (Tree (c a yk y b) (dispatch-on-key key-compare key yk (recur a) y (recur b))))) ) ;; Finds an association with a given key, and deletes it. Returns ;; the (key . value) pair of the deleted association, or #f if it ;; couldn't be found (define (delete-assoc key-compare root key) (define (zip zipper tree) (match (cons zipper tree) ((($ rb-tree#zipper 'Top) . a) tree) ((($ rb-tree#zipper 'Left color xk x b z) . a) (zip z (Tree color a xk x b))) ((($ rb-tree#zipper 'Right color a xk x z) . b) (zip z (Tree color a xk x b))))) ;; bbZip propagates a black deficit up the tree until either ;; the top is reached, or the deficit can be covered. It ;; returns a boolean that is true if there is still a deficit ;; and the zipped tree. (define (bbZip zipper tree) (match (cons zipper tree) ((($ rb-tree#zipper 'Top) . a) (cons #t a)) ;; case 1L ((($ rb-tree#zipper 'Left 'Black xk x ($ rb-tree#tree 'Tree 'Red c yk y d) z) . a) (bbZip (Left R xk x c (Left B yk y d z)) a)) ;; case 3L ((($ rb-tree#zipper 'Left color xk x ($ rb-tree#tree 'Tree 'Black ($ rb-tree#tree 'Tree 'Red c yk y d) wk w e) z) . a) (bbZip (Left color xk x (Tree B c yk y (Tree R d wk w e)) z) a)) ;; case 4L ((($ rb-tree#zipper 'Left color xk x ($ rb-tree#tree 'Tree 'Black c yk y ($ rb-tree#tree 'Tree 'Red d wk w e)) z) . a) (cons #f (zip z (Tree color (Tree B a xk x c) yk y (Tree B d wk w e))))) ;; case 2L ((($ rb-tree#zipper 'Left 'Red xk x ($ rb-tree#tree 'Tree 'Black c yk y d) z) . a) (cons #f (zip z (Tree B a xk x (Tree R c yk y d))))) ;; case 2L ((($ rb-tree#zipper 'Left 'Black xk x ($ rb-tree#tree 'Tree 'Black c yk y d) z) . a) (bbZip z (Tree B a xk x (Tree R c yk y d)))) ;; case 1R ((($ rb-tree#zipper 'Right color ($ rb-tree#tree 'Tree 'Red c yk y d) xk x z) . b) (bbZip (Right R d xk x (Right B c yk y z)) b)) ;; case 3R ((($ rb-tree#zipper 'Right color ($ rb-tree#tree 'Tree 'Black ($ rb-tree#tree 'Tree 'Red c wk w d) yk y e) xk x z) . b) (bbZip (Right color (Tree B c wk w (Tree R d yk y e)) xk x z) b)) ;; case 4R ((($ rb-tree#zipper 'Right color ($ rb-tree#tree 'Tree 'Black c yk y ($ rb-tree#tree 'Tree 'Red d wk w e)) xk x z) . b) (cons #f (zip z (Tree color c yk y (Tree B (Tree R d wk w e) xk x b))))) ;; case 2R ((($ rb-tree#zipper 'Right 'Red ($ rb-tree#tree 'Tree 'Black c yk y d) xk x z) . b) (cons #f (zip z (Tree B (Tree R c yk y d) xk x b)))) ;; case 2R ((($ rb-tree#zipper 'Right 'Black ($ rb-tree#tree 'Tree 'Black c yk y d) xk x z) . b) (bbZip z (Tree B (Tree R c yk y d) xk x b))) (else (cons #f (zip zipper tree))))) (define (delMin tree z) (match tree (($ rb-tree#tree 'Tree 'Red ($ rb-tree#tree 'Empty) yk y b) (values yk y (cons #f (zip z b)))) (($ rb-tree#tree 'Tree 'Black ($ rb-tree#tree Empty) yk y b) (values yk y (bbZip z b))) (($ rb-tree#tree 'Tree color a yk y b) (delMin a (Left color yk y b z))) (($ rb-tree#tree 'Empty) (rb-tree:error 'delete "invalid tree")))) (define join (match-lambda* (( 'Red ($ rb-tree#tree 'Empty) ($ rb-tree#tree 'Empty) z) (zip z (Empty))) (( _ a ($ rb-tree#tree 'Empty) z) (cdr (bbZip z a))) (( _ ($ rb-tree#tree 'Empty) b z) (cdr (bbZip z b))) ((color a b z) (let-values (((xk x b) (delMin b (Top)))) (match b ((#t . b1) (cdr (bbZip z (Tree color a xk x b1)))) ((#f . b1) (zip z (Tree color a xk x b1)))))))) (define (del tree key z) (match tree (($ rb-tree#tree 'Empty) #f) (($ rb-tree#tree 'Tree color a yk y b) (dispatch-on-key key-compare key yk (del a key (Left color yk y b z)) (cons (cons yk y) (join color a b z)) (del b key (Right color a yk y z)))))) (del root key (Top))) (define (get-min root) (define (f root) (match root (($ rb-tree#tree 'Empty) #f) (($ rb-tree#tree 'Tree _ _ ($ rb-tree#tree 'Empty) xk x _) (cons xk x)) (($ rb-tree#tree 'Tree _ a _ _ _) (f a)))) (f root)) (define (get-max root) (define (f root) (match root (($ rb-tree#tree 'Empty) #f) (($ rb-tree#tree 'Tree _ _ xk x ($ rb-tree#tree 'Empty)) (cons xk x)) (($ rb-tree#tree 'Tree _ _ _ _ b) (f b)))) (f root)) (define (for-each-ascending root f) (define (appf f tree) (match tree (($ rb-tree#tree 'Empty) (void)) (($ rb-tree#tree 'Tree _ a k x b) (begin (appf f a) (f (cons k x)) (appf f b))))) (appf f root)) (define (for-each-descending root f) (define (appf f tree) (match tree (($ rb-tree#tree 'Empty) (void)) (($ rb-tree#tree 'Tree _ a k x b) (begin (appf f b) (f (cons k x)) (appf f a))))) (appf f root)) (define (apply-default-clause label key default-clause) (cond ((null? default-clause) (rb-tree:error label "key " key " was not found in the tree")) ((pair? (cdr default-clause)) (rb-tree:error label "default argument must be a single clause")) ((procedure? (car default-clause)) ((car default-clause))) (else (car default-clause)))) ;; functions for walking the tree while keeping a stack of parents to ;; be visited. (define (next lst) (match lst (((and t ($ rb-tree#tree 'Tree _ _ _ _ b)) . rest) (list t (left b rest))) (else (list (Empty) '())))) (define (left t rest) (match t (($ rb-tree#tree 'Empty) rest) ((and t ($ rb-tree#tree 'Tree _ a _ _ _)) (left a (cons t rest))))) (define (start t) (left t '())) ;; Support for constructing red-black trees in linear time from ;; increasing ordered sequences (based on a description by R. Hinze). ;; Note that the elements in the digits are ordered with the largest ;; on the left, whereas the elements of the trees are ordered with the ;; largest on the right. (define-datatype digit digit? (Zero) (One (key identity) (value identity) (tree tree?) (digit digit?)) (Two (key1 identity) (value1 identity) (tree1 tree?) (key2 identity) (value2 identity) (tree2 tree?) (digit digit?))) ;; add an item that is guaranteed to be larger than any in l (define (add-item ak a l) (define (incr ak1 a1 t1 d) (match d (($ digit 'Zero) (One ak1 a1 t1 (Zero))) (($ digit 'One ak2 a2 t2 r) (Two ak1 a1 t1 ak2 a2 t2 r)) (($ digit 'Two ak2 a2 t2 ak3 a3 t3 r) (One ak1 a1 t1 (incr ak2 a2 (Tree B t3 ak3 a3 t2) r))))) (incr ak a (Empty) l)) ;; link the digits into a tree (define (link-all t) (define (link t d) (match d (($ digit 'Zero) t) (($ digit 'One ak a t2 r) (link (Tree B t2 ak a t) r)) (($ digit 'Two ak1 a1 t1 ak2 a2 t2 r) (link (Tree B (Tree R t2 ak2 a2 t1) ak1 a1 t) r)))) (link (Empty) t)) (define (wrap f) (lambda (t1 t2) (match-let (((n result) (f (start t1) (start t2) 0 (Zero)))) (link-all result)))) (define (map-insert t n result) (match t ((($ rb-tree#tree 'Empty) _) (list n result)) ((($ rb-tree#tree 'Tree _ _ xk x _) r) (map-insert (next r) (+ 1 n) (add-item xk x result))))) ;; Creates a map whose domain is the union of the domains of the two ;; input maps, using the supplied function to define the map on ;; elements that are in both domains. (define (union-with merge-fn) (define (union k1 k2 key-compare) (lambda (t1 t2 n result) (let recur ((t1 t1) (t2 t2) (n n) (result result)) (match (list (next t1) (next t2)) (((($ rb-tree#tree 'Empty) _) (($ rb-tree#tree 'Empty) _)) (list n result)) (((($ rb-tree#tree 'Empty) _) t2) (map-insert t2 n result)) ((t1 (($ rb-tree#tree 'Empty) _)) (map-insert t1 n result)) (((($ rb-tree#tree 'Tree _ _ xk x _) r1) (($ rb-tree#tree 'Tree _ _ yk y _) r2)) (let ((xk1 (k1 xk)) (yk1 (k2 yk))) (let ((c (key-compare xk1 yk1))) (cond ((negative? c) (recur r1 t2 (+ 1 n) (add-item xk1 x result))) ((zero? c) (recur r1 r2 (+ 1 n) (add-item xk1 (merge-fn x y) result))) ((positive? c) (recur t1 r2 (+ 1 n) (add-item yk1 y result))))))) )))) (wrap union)) (define (union-withi merge-fn) (define (union k1 k2 key-compare) (lambda (t1 t2 n result) (let recur ((t1 t1) (t2 t2) (n n) (result result)) (match (list (next t1) (next t2)) (((($ rb-tree#tree 'Empty) _) (($ rb-tree#tree 'Empty) _)) (list n result)) (((($ rb-tree#tree 'Empty) _) t2) (map-insert t2 n result)) ((t1 (($ rb-tree#tree 'Empty) _)) (map-insert t1 n result)) (((($ rb-tree#tree 'Tree _ _ xk x _) r1) (($ rb-tree#tree 'Tree_ _ yk y _) r2)) (let ((xk1 (k1 xk)) (yk1 (k2 yk))) (let ((c (key-compare xk1 yk1))) (cond ((negative? c) (recur r1 t2 (+ 1 n) (add-item xk1 x result))) ((zero? c) (recur r1 r2 (+ 1 n) (add-item xk1 (merge-fn xk1 x y) result))) ((positive? c) (recur t1 r2 (+ 1 n) (add-item yk1 y result))))))) )))) (wrap union)) (define (get-depth root) (let loop ((node root) (level 0)) (match node (($ rb-tree#tree 'Empty) level) (($ rb-tree#tree 'Tree _ a _ _ b) (max (loop a (+ 1 level)) (loop b (+ 1 level))))))) (define (get-size root) (let loop ((node root)) (match node (($ rb-tree#tree 'Empty) 0) (($ rb-tree#tree 'Tree _ a _ _ b) (+ 1 (loop a) (loop b))) )) ) (define *eof-object* (read (open-input-string ""))) (define (eof-object) *eof-object*) (define (value-generator root) (let ((s (match (start root) ((fst . rst) (make-parameter (list fst rst))) (else (error 'value-generator "empty root"))))) (lambda () (let* ((sval (s)) (t (car sval)) (rst (cadr sval))) (s (next rst)) (cases tree t (Empty () (eof-object)) (Tree (c l k v r) v)) )) )) (define (key-generator root) (let ((s (match (start root) ((fst . rst) (make-parameter (list fst rst))) (else (error 'value-generator "empty root"))))) (lambda () (let* ((sval (s)) (t (car sval)) (rst (cadr sval))) (s (next rst)) (cases tree t (Empty () (eof-object)) (Tree (c l k v r) k)) )) )) (define (rb-tree-map key-compare #!key (new-root (Empty))) (let ( (root new-root) ) (object ;; persistent map behaviors ((persistent-map? self) #t) ((size self) (get-size root)) ((empty? self) (cases tree root (Empty () #t) (else #f))) ((get self key) (or (find-assoc key-compare root key) (apply-default-clause 'get key '()))) ((get/default self key default-clause) (or (find-assoc key-compare root key) (apply-default-clause 'get key (list default-clause)))) ((get-value self key) (or (find-ref key-compare root key) (apply-default-clause 'get-value key '()))) ((get-value/default self key default-clause) (or (find-ref key-compare root key) (apply-default-clause 'get-value key (list default-clause)))) ((get-min self) (get-min root)) ((get-max self) (get-max root)) ((generate self p f g seed) (let recur ((new-root root) (seed seed)) (if (p seed) (rb-tree-map key-compare new-root: new-root) (let* ((item (f seed)) (new-root (insert key-compare (lambda (x ax) x) new-root (car item) (cdr item)))) (recur new-root (g seed)))))) ((put self key value) (let ((new-root (insert key-compare (lambda (x ax) x) root key value))) (rb-tree-map key-compare new-root: new-root))) ((update self key value merge-fn) (let ((new-root (insert key-compare merge-fn root key value))) (rb-tree-map key-compare new-root: new-root))) ((delete self key) (or (let* ((item+tree (delete-assoc key-compare root key)) (item (car item+tree)) (new-root (cdr item+tree))) (cons item (rb-tree-map key-compare new-root: new-root))) (apply-default-clause 'delete key '()))) ((for-each-ascending self f) (for-each-ascending root f)) ((for-each-descending self f) (for-each-descending root f)) ;; collection behaviors ((collection? self) #t) ((gen-keys self) (key-generator root)) ((gen-elts self) (value-generator root)) ((for-each-key self proc) (for-each-ascending root (lambda (item) (proc (car item))))) ((for-each-elt self proc) (for-each-ascending root (lambda (item) (proc (cdr item))))) )) ) )