; avl trees
(define (tree k v l r)
(vector k v l r (+ (max (ht l) (ht r)) 1)
(+ (size l) (size r) 1)))
(define (key t) (vector-ref t 0))
(define (val t) (vector-ref t 1))
(define (lkid t) (vector-ref t 2))
(define (rkid t) (vector-ref t 3))
(define (ht t) (vector-ref t 4))
(define (size t) (vector-ref t 5))
(define (bal t) (- (ht (lkid t)) (ht (rkid t))))
(define nil (vector 'nil 'nil 'nil 'nil 0 0))
(vector-set! nil 2 nil)
(vector-set! nil 3 nil)
(define (nil? t) (eq? t nil))
(define (rot-left t)
(if (nil? t) t
(tree (key (rkid t))
(val (rkid t))
(tree (key t) (val t) (lkid t) (lkid (rkid t)))
(rkid (rkid t)))))
(define (rot-right t)
(if (nil? t) t
(tree (key (lkid t))
(val (lkid t))
(lkid (lkid t))
(tree (key t) (val t) (rkid (lkid t)) (rkid t)))))
(define (balance t)
(let ((b (bal t)))
(cond ((< (abs b) 2) t)
((positive? b)
(if (< -1 (bal (lkid t))) (rot-right t)
(rot-right (tree (key t) (val t)
(rot-left (lkid t)) (rkid t)))))
((negative? b)
(if (< (bal (rkid t)) 1) (rot-left t)
(rot-left (tree (key t) (val t)
(lkid t) (rot-right (rkid t)))))))))
(define (lookup lt? t k)
(cond ((nil? t) #f)
((lt? k (key t)) (lookup lt? (lkid t) k))
((lt? (key t) k) (lookup lt? (rkid t) k))
(else (cons k (val t)))))
(define (insert lt? t k v)
(cond ((nil? t) (tree k v nil nil))
((lt? k (key t))
(balance (tree (key t) (val t)
(insert lt? (lkid t) k v) (rkid t))))
((lt? (key t) k)
(balance (tree (key t) (val t)
(lkid t) (insert lt? (rkid t) k v))))
(else (tree k v (lkid t) (rkid t)))))
(define (delete-successor t)
(if (nil? (lkid t)) (values (rkid t) (key t) (val t))
(call-with-values
(lambda () (delete-successor (lkid t)))
(lambda (l k v)
(values (balance (tree (key t) (val t) l (rkid t))) k v)))))
(define (delete lt? t k)
(cond ((nil? t) nil)
((lt? k (key t))
(balance (tree (key t) (val t)
(delete lt? (lkid t) k) (rkid t))))
((lt? (key t) k)
(balance (tree (key t) (val t)
(lkid t) (delete lt? (rkid t) k))))
((nil? (lkid t)) (rkid t))
((nil? (rkid t)) (lkid t))
(else (call-with-values
(lambda () (delete-successor (rkid t)))
(lambda (r k v) (balance (tree k v (lkid t) r)))))))
(define (enlist t)
(cond ((nil? t) (list))
((and (nil? (lkid t)) (nil? (rkid t)))
(list (cons (key t) (val t))))
(else (append (enlist (lkid t))
(list (cons (key t) (val t)))
(enlist (rkid t))))))
(define (nth t n)
(if (negative? n) (error 'nth "must be non-negative")
(let ((s (size (lkid t))))
(cond ((< n s) (nth (lkid t) n))
((< s n) (nth (rkid t) (- n s 1)))
((nil? t) #f)
(else (cons (key t) (val t)))))))
(define (rank lt? t k)
(let loop ((t t) (s (size (lkid t))))
(cond ((nil? t) #f)
((lt? k (key t))
(loop (lkid t) (size (lkid (lkid t)))))
((lt? (key t) k)
(loop (rkid t) (+ s (size (lkid (rkid t))) 1)))
(else s))))
(define (avl-map proc t) ; (proc key value)
(if (nil? t) nil
(tree (key t) (proc (key t) (val t))
(avl-map proc (lkid t))
(avl-map proc (rkid t)))))
(define (avl-fold proc base t) ; (proc key value base)
(if (nil? t) base
(avl-fold proc
(proc (key t) (val t)
(avl-fold proc base (lkid t)))
(rkid t))))
(define (avl-for-each proc t) ; (proc key value)
(unless (nil? t)
(avl-for-each proc (lkid t))
(proc (key t) (val t))
(avl-for-each proc (rkid t))))
(define (list->avl lt? xs)
(let loop ((xs xs) (t nil))
(if (null? xs) t
(loop (cdr xs) (insert lt? t (caar xs) (cdar xs))))))
(avl-for-each (lambda (k v) (display k) (display " ") (display v) (newline))
(avl-map (lambda (k v) (+ v v))
(list->avl <
(map (lambda (x) (cons x x))
(list 0 1 2 3 4 5 6 7 8 9)))))
(display
(avl-fold (lambda (k v base) (+ v base)) 0
(list->avl <
(map (lambda (x) (cons x x))
(list 0 1 2 3 4 5 6 7 8 9)))))