502 lines
12 KiB
Scheme
502 lines
12 KiB
Scheme
|
|
(declare (unit util-bst))
|
|
|
|
(import duck)
|
|
|
|
(module*
|
|
util-bst
|
|
#:doc ("Binary Search Tree implementation")
|
|
(
|
|
make-bst
|
|
bst?
|
|
|
|
bst-EQ?
|
|
|
|
bst-count
|
|
|
|
bst-empty?
|
|
bst-ref
|
|
bst-contains?
|
|
|
|
bst-set
|
|
bst-remove
|
|
|
|
bst-keys
|
|
|
|
bst->kvv
|
|
kvv->bst
|
|
kvv-filter
|
|
|
|
bst-balance
|
|
|
|
bst-find-pair
|
|
bst-filter-pairs
|
|
|
|
bst-map-list
|
|
bst-map-bst
|
|
|
|
list->bst
|
|
|
|
bst-update
|
|
|
|
bst-filter
|
|
bst-reduce
|
|
|
|
bst-equal?
|
|
|
|
util-bst-tests!
|
|
)
|
|
|
|
(import scheme
|
|
(chicken condition)
|
|
(chicken sort)
|
|
util-tag
|
|
testing
|
|
racket-kwargs)
|
|
|
|
;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;
|
|
;; Node
|
|
|
|
;; Creates BST node with no children
|
|
(define (make-bst-node key value)
|
|
(cons (cons key value)
|
|
(cons #f #f)))
|
|
|
|
;; Read-only accessors to BST node
|
|
(define bst-node-kv car)
|
|
(define bst-node-key caar)
|
|
(define bst-node-value cdar)
|
|
(define bst-node-left cadr)
|
|
(define bst-node-right cddr)
|
|
|
|
;; Returns BST node with updated node value
|
|
(define (set-bst-node-value n v)
|
|
(cons (cons (bst-node-key n) v)
|
|
(cdr n)))
|
|
|
|
;; Updates BST node left child
|
|
(define (set-bst-node-left n l)
|
|
(cons (car n)
|
|
(cons l (bst-node-right n))))
|
|
|
|
;; Updates BST node right child
|
|
(define (set-bst-node-right n r)
|
|
(cons (car n)
|
|
(cons (bst-node-left n) r)))
|
|
|
|
;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;
|
|
;; BST
|
|
|
|
;; Unique tag
|
|
(define TAG-BST (make-tag bst))
|
|
|
|
(define/doc (make-bst subtag EQ? <?)
|
|
("Creates empty BST with given comparators")
|
|
(cons TAG-BST
|
|
(cons (cons #f 0)
|
|
(cons subtag
|
|
(cons EQ? <?)))))
|
|
|
|
;; BST accessors
|
|
(define bst-tag car)
|
|
(define bst-root+count cadr)
|
|
(define bst-root caadr)
|
|
(define bst-count cdadr)
|
|
(define bst-type cddr)
|
|
(define bst-subtag caddr)
|
|
(define bst-comparators cdddr)
|
|
(define bst-EQ? cadddr)
|
|
(define bst-<? cddddr)
|
|
|
|
;; Update BST root node
|
|
(define (set-bst-root bst root)
|
|
(cons (car bst)
|
|
(cons (cons root
|
|
(bst-count bst))
|
|
(cddr bst))))
|
|
|
|
;; Update BST count
|
|
(define (set-bst-count bst count)
|
|
(cons (car bst)
|
|
(cons (cons (bst-root bst)
|
|
count)
|
|
(cddr bst))))
|
|
|
|
;; Update BST root node and count
|
|
(define (set-bst-root+count bst root count)
|
|
(cons (car bst)
|
|
(cons (cons root count)
|
|
(cddr bst))))
|
|
|
|
(define/doc ((bst? subtag) v)
|
|
("Curried predicate for particular bst type.")
|
|
(and (pair? v)
|
|
(eq? (bst-tag v) TAG-BST)
|
|
(eq? (bst-subtag v) subtag)))
|
|
|
|
(define/doc (bst-empty? bst)
|
|
("Returns #t if given BST is empty.")
|
|
(not (bst-root bst)))
|
|
|
|
;; Wrapper to setup comparators
|
|
(define-syntax let-comparators
|
|
(syntax-rules ()
|
|
((_ (EQ? <? bst) expr ...)
|
|
(let ((EQ? (bst-EQ? bst))
|
|
(<? (bst-<? bst)))
|
|
expr ...))))
|
|
|
|
(define/doc (bst-ref bst k . vs)
|
|
("Retrieves value associated with given key.")
|
|
(let-comparators
|
|
(EQ? <? bst)
|
|
(let loop ((n (bst-root bst)))
|
|
(if n
|
|
(let ((nk (bst-node-key n)))
|
|
(if (EQ? k nk)
|
|
(bst-node-value n)
|
|
(loop (if (<? k nk)
|
|
(bst-node-left n)
|
|
(bst-node-right n)))))
|
|
(if (null? vs)
|
|
(error 'bst-ref "Key does not exist" k)
|
|
(car vs))))))
|
|
|
|
(define/doc (bst-contains? bst k)
|
|
("Predicate for key existence in BST.")
|
|
(handle-exceptions
|
|
ex
|
|
#f
|
|
(let ()
|
|
(bst-ref bst k)
|
|
#t)))
|
|
|
|
(define/doc (bst-set bst k v)
|
|
("Sets given key to given value and updates count if needed.")
|
|
(let-comparators
|
|
(EQ? <? bst)
|
|
(let-values (((new-root count-add)
|
|
(let loop ((n (bst-root bst)))
|
|
(if n
|
|
(let ((nk (bst-node-key n)))
|
|
(if (EQ? k nk)
|
|
(values (set-bst-node-value n v) 0)
|
|
(if (<? k nk)
|
|
(let-values (((new-left-node count-add)
|
|
(loop (bst-node-left n))))
|
|
(values (set-bst-node-left n new-left-node)
|
|
count-add))
|
|
(let-values (((new-right-node count-add)
|
|
(loop (bst-node-right n))))
|
|
(values (set-bst-node-right n new-right-node)
|
|
count-add)))))
|
|
(values (make-bst-node k v) 1)))))
|
|
(set-bst-root+count bst
|
|
new-root
|
|
(+ (bst-count bst) count-add)))))
|
|
|
|
(define/doc (bst-remove bst k . nos)
|
|
("Removes given key from the BST.")
|
|
(let-comparators
|
|
(EQ? <? bst)
|
|
(let-values
|
|
(((new-root sub-count)
|
|
(let loop ((n (bst-root bst)))
|
|
(if n
|
|
(let ((nk (bst-node-key n)))
|
|
(if (EQ? nk k)
|
|
(values (if (bst-node-left n)
|
|
(if (bst-node-right n)
|
|
(let aloop ((an (bst-node-right n)))
|
|
(if an
|
|
(set-bst-node-left an
|
|
(aloop (bst-node-left an)))
|
|
(bst-node-left n)))
|
|
(bst-node-left n))
|
|
(if (bst-node-right n)
|
|
(bst-node-right n)
|
|
#f))
|
|
1)
|
|
(if (<? k nk)
|
|
(let-values (((new-left-node sub-count)
|
|
(loop (bst-node-left n))))
|
|
(values (set-bst-node-left n new-left-node)
|
|
sub-count))
|
|
(let-values (((new-right-node sub-count)
|
|
(loop (bst-node-right n))))
|
|
(values (set-bst-node-right n new-right-node)
|
|
sub-count)))))
|
|
(if (and (not (null? nos))
|
|
(car nos))
|
|
(values #f 0)
|
|
(error 'bst-remove "Key does not exist" k))))))
|
|
(if (eq? sub-count 0)
|
|
bst
|
|
(set-bst-root+count bst
|
|
new-root
|
|
(- (bst-count bst) sub-count))))))
|
|
|
|
;; Unique tag for static vector representation that retains type info
|
|
(define TAG-KVV (make-tag kvv))
|
|
|
|
;; Iterates over given BST and passes the KV pairs to given proc
|
|
(define (bst-iter-kv bst proc)
|
|
(let-comparators
|
|
(EQ? <? bst)
|
|
(let loop ((n (bst-root bst)))
|
|
(when n
|
|
(loop (bst-node-left n))
|
|
(proc (bst-node-kv n))
|
|
(loop (bst-node-right n))))))
|
|
|
|
;; Used for iterating over multiple BSTs simultaneously
|
|
(define (bst-kv-iterator bst)
|
|
(let* ((break #f)
|
|
(resume #f)
|
|
(yield (lambda (val)
|
|
(call/cc
|
|
(lambda (r)
|
|
(set! resume r)
|
|
(break val))))))
|
|
(lambda ()
|
|
(call/cc
|
|
(lambda (cc)
|
|
(set! break cc)
|
|
(if resume
|
|
(resume '())
|
|
(bst-iter-kv bst yield))
|
|
#f)))))
|
|
|
|
(define/doc (bst-keys bst)
|
|
("Returns all the keys contained in given dictionary.")
|
|
(let ((res '()))
|
|
(bst-iter-kv bst
|
|
(lambda (kv)
|
|
(set! res (cons (car kv) res))))
|
|
(reverse res)))
|
|
|
|
;; Converts BST to KVV
|
|
(define (bst->kvv bst)
|
|
(let ((vec (make-vector (bst-count bst)))
|
|
(idx 0))
|
|
(bst-iter-kv bst
|
|
(lambda (kv)
|
|
(vector-set! vec idx kv)
|
|
(set! idx (add1 idx))))
|
|
(cons TAG-KVV
|
|
(cons vec
|
|
(cddr bst)))))
|
|
|
|
;; Converts a vector of key-value cons cells to BST dictionary
|
|
(define (kvv->bst kvv)
|
|
(let* ((vec (cadr kvv))
|
|
(count (vector-length vec)))
|
|
(cons TAG-BST
|
|
(cons (cons (let loop ((s 0)
|
|
(e count))
|
|
(if (= s e)
|
|
#f
|
|
(let* ((c (- e s))
|
|
(h (quotient c 2))
|
|
(m (+ s h))
|
|
(kv (vector-ref vec m)))
|
|
(cons kv
|
|
(cons (loop s m)
|
|
(loop (add1 m) e))))))
|
|
count)
|
|
(cddr kvv)))))
|
|
|
|
;; Returns KVV that contains only KV pairs matching pred?
|
|
(define (kvv-filter kvv pred?)
|
|
(cons (car kvv)
|
|
(cons (apply vector
|
|
(let ((vec (cadr kvv)))
|
|
(let loop ((idx (sub1 (vector-length vec)))
|
|
(res '()))
|
|
(if (< idx 0)
|
|
res
|
|
(let ((kv (vector-ref vec idx)))
|
|
(loop (if (pred? (car kv) (cdr kv))
|
|
(cons kv res)
|
|
res)
|
|
(sub1 idx)))))))
|
|
(cddr kvv))))
|
|
|
|
(define/doc (bst-balance bst)
|
|
("Unconditionally balances the BST.")
|
|
(kvv->bst
|
|
(bst->kvv bst)))
|
|
|
|
(define/doc (bst-find-pair bst pred?)
|
|
("Finds pair by predicate that accepts both key and value.")
|
|
(call/cc
|
|
(lambda (cc)
|
|
(bst-iter-kv bst
|
|
(lambda (kv)
|
|
(when (pred? (car kv) (cdr kv))
|
|
(cc kv))))
|
|
#f)))
|
|
|
|
(define/doc (bst-filter-pairs bst pred?)
|
|
("Returns a list of key-value pairs matching predicate.")
|
|
(let ((res '()))
|
|
(bst-iter-kv bst
|
|
(lambda (kv)
|
|
(when (pred? (car kv) (cdr kv))
|
|
(set! res (cons kv res)))))
|
|
(reverse res)))
|
|
|
|
(define/doc (bst-map-list bst proc)
|
|
("Returns arbitrary list created by mapping all elements.")
|
|
(let ((res '()))
|
|
(bst-iter-kv bst
|
|
(lambda (kv)
|
|
(set! res
|
|
(cons res
|
|
(proc (car kv) (cdr kv))))))
|
|
(reverse res)))
|
|
|
|
(define/doc (bst-map-bst bst proc)
|
|
("Returns a new dictionary with all values processed (keys are left intact).")
|
|
(set-bst-root bst
|
|
(let loop ((n (bst-root bst)))
|
|
(if n
|
|
(cons (cons (caar n) (proc (caar n) (cdar n)))
|
|
(cons (loop (cadr n))
|
|
(loop (cddr n))))
|
|
#f))))
|
|
|
|
(define/doc (list->bst lst subtag EQ? <?)
|
|
("Converts list of pairs into BST dictionary.")
|
|
(kvv->bst
|
|
(cons TAG-KVV
|
|
(cons (list->vector
|
|
(sort lst (lambda (akv bkv)
|
|
(<? (car akv) (car bkv)))))
|
|
(cons subtag
|
|
(cons EQ? <?))))))
|
|
|
|
(define*/doc (bst-update bst k proc (v #f))
|
|
("Functional update with optional default value (defaults to #f).")
|
|
(let-comparators
|
|
(EQ? <? bst)
|
|
(let-values (((new-root add-count)
|
|
(let loop ((n (bst-root bst)))
|
|
(if n
|
|
(let ((nk (bst-node-key n)))
|
|
(if (EQ? k nk)
|
|
(values (set-bst-node-value n (proc (bst-node-value n)))
|
|
0)
|
|
(if (<? k nk)
|
|
(let-values (((new-left-node add-count)
|
|
(loop (bst-node-left n))))
|
|
(values (set-bst-node-left n new-left-node)
|
|
add-count))
|
|
(let-values (((new-right-node add-count)
|
|
(loop (bst-node-right n))))
|
|
(values (set-bst-node-right n new-right-node)
|
|
add-count)))))
|
|
(values (make-bst-node k (proc v)) 1)))))
|
|
(set-bst-root+count bst
|
|
new-root
|
|
(+ (bst-count bst) add-count)))))
|
|
|
|
(define/doc (bst-filter bst pred?)
|
|
("Returns a BST with only KV pairs matching the predicate which must
|
|
accept two arguments.")
|
|
(kvv->bst
|
|
(cons TAG-KVV
|
|
(cons (apply vector (bst-filter-pairs bst pred?))
|
|
(cddr bst)))))
|
|
|
|
(define/doc (bst-reduce bst proc init)
|
|
("Like generic reduce, the proc gets accumulator, key and value
|
|
arguments.")
|
|
(let ((acc init))
|
|
(bst-iter-kv bst
|
|
(lambda (kv)
|
|
(set! acc (proc acc (car kv) (cdr kv)))))
|
|
acc))
|
|
|
|
;; Returns true if these are compatible BSTs
|
|
(define (bst-compat? b1 b2)
|
|
(and (eq? (bst-tag b1)
|
|
(bst-tag b2))
|
|
(eq? (bst-subtag b1)
|
|
(bst-subtag b2))
|
|
(eq? (bst-EQ? b1)
|
|
(bst-EQ? b2))
|
|
(eq? (bst-<? b1)
|
|
(bst-<? b2))))
|
|
|
|
(define*/doc (bst-equal? b1 b2 (equality? equal?))
|
|
("Returns true if both BSTs contain the same keys and values.")
|
|
(if (bst-compat? b1 b2)
|
|
(let-comparators
|
|
(EQ? _ b1)
|
|
(let ((g1 (bst-kv-iterator b1))
|
|
(g2 (bst-kv-iterator b2)))
|
|
(let loop ()
|
|
(let ((kv1 (g1))
|
|
(kv2 (g2)))
|
|
(if (and (not kv1)
|
|
(not kv2))
|
|
#t
|
|
(if (and (EQ? (car kv1)
|
|
(car kv2))
|
|
(equality? (cdr kv1)
|
|
(cdr kv2)))
|
|
(loop)
|
|
#f))))))
|
|
#f))
|
|
|
|
;; Module self-tests
|
|
(define (util-bst-tests!)
|
|
(run-tests
|
|
util-bst
|
|
(test-equal? make-bst
|
|
(make-bst 'fixnum eq? <)
|
|
`(,TAG-BST
|
|
. ((#f . 0)
|
|
. (fixnum . (,eq? . ,<)))))
|
|
(test-true bst?
|
|
((bst? 'fixnum) (make-bst 'fixnum eq? <)))
|
|
(test-false bst?
|
|
((bst? 'fixnum) (make-bst 'string eq? <)))
|
|
(test-false bst?
|
|
((bst? 'fixnum) "string"))
|
|
(test-true bst-empty?
|
|
(bst-empty? (make-bst 'fixnum eq? <)))
|
|
(test-true bst-contains?
|
|
(bst-contains?
|
|
(bst-set
|
|
(make-bst 'fixnum eq? <)
|
|
1 2) 1))
|
|
(test-false bst-contains?
|
|
(bst-contains? (make-bst 'fixnum eq? <) 1))
|
|
(test-equal? bst-count
|
|
(bst-count
|
|
(make-bst 'fixnum eq? <))
|
|
0)
|
|
(test-equal? bst-set
|
|
(bst-count
|
|
(bst-set
|
|
(make-bst 'fixnum eq? <)
|
|
1 2))
|
|
1)
|
|
(test-true bst-remove
|
|
(bst-empty?
|
|
(bst-remove
|
|
(bst-set
|
|
(make-bst 'fixnum eq? <)
|
|
1 2) 1)))
|
|
(test-equal? bst-balance
|
|
(bst-count
|
|
(bst-balance
|
|
(bst-set
|
|
(make-bst 'fixnum eq? <)
|
|
1 2)))
|
|
1)
|
|
))
|
|
|
|
)
|