hackerbase/src/util-bst.scm

214 lines
4.6 KiB
Scheme

(declare (unit util-bst))
(import duck)
(module*
util-bst
#:doc ("Binary Search Tree implementation")
(
make-bst
bst?
bst-empty?
bst-ref
bst-contains?
bst-set
;;bst-remove ;;
;;bst->kvv
;;kvv->bst
;;kvv-filter
;;bst-balance
util-bst-tests!
)
(import scheme
(chicken condition)
util-tag
testing)
;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;
;; 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)))))
;; 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-count
(bst-count
(bst-set
(make-bst 'fixnum eq? <)
1 2))
1)
))
)
(import util-bst)
(util-bst-tests!)