#lang scheme (require (planet schematics/schemeunit)) (provide (struct-out tree) (struct-out leaf) (struct-out node) tree-map tree-map/n build-tree tr:make-tree) (define-struct tree (val) #:prefab) (define-struct (leaf tree) () #:prefab) (define-struct (node tree) (left right) #:prefab) ;; Not currently used, but might come in handy at some point. ;; [Tree X] -> [Seq X] (define (tree->preorder-seq t) (make-do-sequence ;; position is a forest [Listof [Tree X]] (lambda () (values (lambda (x) (tree-val (car x))) (lambda (p) (cond [(leaf? (car p)) (cdr p)] [else (cons (node-left (car p)) (cons (node-right (car p)) (cdr p)))])) (list t) cons? void void)))) ;; [X -> Y] [Tree X] -> [Tree Y] (define (tree-map f t) (cond [(leaf? t) (make-leaf (f (tree-val t)))] [else (make-node (f (tree-val t)) (tree-map f (node-left t)) (tree-map f (node-right t)))])) ;; [X Y Z ... -> R] [List [Tree X] [Tree Y] [Tree Z] ...] -> [Tree R] (define (tree-map/n f ts) (let recr ((ts ts)) (cond [(leaf? (car ts)) (make-leaf (apply f (map tree-val ts)))] [else (make-node (apply f (map tree-val ts)) (recr (map node-left ts)) (recr (map node-right ts)))]))) ;; Consumes n = 2^i-1 and produces 2^(i-1)-1. ;; Nat -> Nat (define (half n) (arithmetic-shift n -1)) ;; Nat [Nat -> X] -> [Tree X] ;; like build-list, but for complete binary trees (define (build-tree i f) ;; i = 2^j-1 (let rec ((i i) (o 0)) (cond [(= 1 i) (make-leaf (f o))] [else (make-node (f o) (rec (half i) (+ o 1)) (rec (half i) (+ o 1 (half i))))]))) ;; Nat X -> [Tree X] (define (tr:make-tree i x) ;; i = 2^j-1 (let rec ((i i)) (cond [(= 1 i) (make-leaf x)] [else (make-node x (rec (half i)) (rec (half i)))]))) ;; --------------------------------------------------------------------------- ;; Test suite (define/provide-test-suite tree-tests (check-equal? (tree-map add1 (make-leaf 0)) (make-leaf 1)) (check-equal? (tree-map add1 (make-node 0 (make-leaf 1) (make-leaf 2))) (make-node 1 (make-leaf 2) (make-leaf 3))) (check-equal? (tree-map/n cons (list (make-leaf 'a) (make-leaf 'z))) (make-leaf '(a . z))) (check-equal? (tree-map/n cons (list (make-node 'a (make-leaf 'b) (make-leaf 'c)) (make-node 'z (make-leaf 'y) (make-leaf 'x)))) (make-node '(a . z) (make-leaf '(b . y)) (make-leaf '(c . x)))) (check-equal? (build-tree 1 (lambda (i) i)) (make-leaf 0)) (check-equal? (build-tree 3 (lambda (i) i)) (make-node 0 (make-leaf 1) (make-leaf 2))) (check-equal? (build-tree 7 (lambda (i) i)) (make-node 0 (make-node 1 (make-leaf 2) (make-leaf 3)) (make-node 4 (make-leaf 5) (make-leaf 6)))) (check-equal? (build-tree 1 (lambda (i) 'x)) (make-leaf 'x)) (check-equal? (build-tree 3 (lambda (i) 'x)) (make-node 'x (make-leaf 'x) (make-leaf 'x))) (check-equal? (for/list ([i (tree->preorder-seq (make-leaf 0))]) i) (list 0)) (check-equal? (for/list ([i (tree->preorder-seq (make-node 0 (make-node 1 (make-leaf 2) (make-leaf 3)) (make-node 4 (make-leaf 5) (make-leaf 6))))]) i) (list 0 1 2 3 4 5 6)))