## AVL Trees, Extended

### November 29, 2011

It is easy to add a size field to each node in the tree that stores the number of non-nil nodes in the tree (including the current node); we have to modify both constructors and provide a selector:

```(define (tree k v l r)   (vector k v l r (+ (max (ht l) (ht r)) 1)                   (+ (size l) (size r) 1)))```

`(define nil (vector 'nil 'nil 'nil 'nil 0 0))`

`(define (size t) (vector-ref t 5))`

Once we have the size of each tree, we can write the `nth` and `rank` functions. We choose to say that the first element of the tree is at index 0, though if you prefer to start the index at 1 that’s not hard to arrange:

```(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))))```

The key to both these functions is that the index of a node is the size of its left child. As long as we are descending a tree via its left child, we just take the index of the node. But when we have to move down the tree via its right child, we have to adjust the size to include the half of the tree that we have excluded. Thus, unlike the `insert` and `delete` operations, the code for the `nth` and `rank` operations has different operations on the left and right subtrees.

The three common operations on lists — `map`, `fold` and `for-each` — translate neatly to AVL trees.

```(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))))```

The final function is `list->avl` which inserts the elements of a list into an AVL tree:

```(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))))))```

Here are some examples:

```> (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))))) 0 0 1 2 2 4 3 6 4 8 5 10 6 12 7 14 8 16 9 18 > (avl-fold (lambda (k v base) (+ v base)) 0     (list->avl < (map (lambda (x) (cons x x)) (range 10)))) 45```

You can run the program at http://programmingpraxis.codepad.org/pwYpvDFw.