0. #lang racket
1. (require racket/pretty)
2. (provide insert lookup btree)
3.  
4. (struct Btree (tree = > <) #:transparent)
5. (struct Empty ())
6. (struct Node (key value size left right) #:transparent)
7.  
8. (define (btree = > <)
9.   (Btree (Empty) = > <))
10.  
11. (define (node-size n)
12.   (match n
13.     [(Empty) 0]
14.     [(Node _ _ s _ _) s]))
15.  
16. (define (insert-aux tree k v = > <)
17.   (match tree
18.     [(Node tk tvs ts tl tr)
19.      (cond
20.        [(= tk k) (Node tk (set-add tvs v) ts tl tr)]
21.        [(> tk k) (let* ([new-left (insert-aux tl k v = > <)]
22.                         [new-size (+ 1 (node-size new-left) (node-size tr))])
23.                    (Node tk tvs new-size new-left tr))]
24.        [(< tk k) (let* ([new-right (insert-aux tr k v = > <)]
25.                         [new-size (+ 1 (node-size tl) (node-size new-right))])
26.                    (Node tk tvs new-size tl new-right))])]
27.     [(Empty) (Node k (set v) 1 (Empty) (Empty))]))
28.  
29. (define (insert btree k v)
30.   (match-let ([(Btree tree = > <) btree])
31.     (Btree (insert-aux tree k v = > <) = > <)))
32.  
33. (define (insert*  btree . inits)
34.   (foldl (lambda (x m) (insert m (first x) (second x))) btree inits))
35.  
36.  
37. (define (lookup-aux tree k = > <)
38.   (match tree
39.     [(Node tk tvs ts tl tr)
40.      (cond
41.        [(= tk k) tvs]
42.        [(> tk k) (lookup-aux tl k = > <)]
43.        [(< tk k) (lookup-aux tr k = > <)])]
44.     [(Empty) (set)]))
45.  
46. (define (lookup btree k)
47.   (match-let ([(Btree tree = > <) btree])
48.     (lookup-aux tree k = > <)))
49.  
50. (define (yield-tree yield tree)
51.   (match tree
52.     [(Node tk tvs ts tl tr)
53.      (yield-tree yield tl)
54.      (yield (list tk tvs))
55.      (yield-tree yield tr)]
56.     [(Empty) (void)]))
57.  
58. (define (btree->vector bt)
59.   (let ([i 0]
60.         [buffer (make-vector (node-size (Btree-tree bt)))])
61.     (yield-tree
62.      (lambda (x)
63.        (vector-set! buffer i x)
64.        (set! i (add1 i))
65.        i)
66.      (Btree-tree bt))
67.     buffer))
68.  
69. (define (btree->alist bt)
70.   (vector->list (btree->vector bt)))
71.  
72. (define (alist->btree alist bt)
73.   (match-let ([(Btree tree = > <) bt])
74.       (foldl (lambda (pair tree)
75.                (insert tree (car pair) (cadr pair)))
76.              bt
77.              alist)))
78.  
79. (define (partial-lookup-aux tree k = > <)
80.   (match tree
81.     [(Node tk tvs ts tl tr)
82.      `(cond
83.          [(,= ,tk ,k) ,(cons 'list (set->list tvs))]
84.          [(,> ,tk ,k) ,(partial-lookup-aux tl k = > <)]
85.          [(,< ,tk ,k) ,(partial-lookup-aux tr k = > <)])]
86.     [(Empty) '(list)]))
87.  
88. (define (partial-lookup btree k)
89.   (match-let ([(Btree tree = > <) btree])
90.     ;; procedures need quoted reprs
91.     (let ([= '=] [> '>] [< '<])
92.       (partial-lookup-aux tree k = > <))))
93.  
94.  
95. (define test-cases '((1 'a) (2 'b) (3 'c) (1 'd)))
96. (define tree1 (alist->btree test-cases (btree = > <)))
97. (define pevald #`(lambda (x) #,(partial-lookup tree1 #'x)))
98. (define runnable (eval pevald))
99. (pretty-print (syntax->datum pevald))
100. #| =>
101. '(lambda (x)
102.    (cond
103.     ((= 1 x) (list 'a 'd))
104.     ((> 1 x) (list))
105.     ((< 1 x)
106.      (cond
107.       ((= 2 x) (list 'b))
108.       ((> 2 x) (list))
109.       ((< 2 x)
110.        (cond ((= 3 x) (list 'c)) ((> 3 x) (list)) ((< 3 x) (list))))))))
111. |#
112.  
113. (map runnable '(1 2 3))
114. ;; => '((a d) (b) (c))

=>