133 lines
7.8 KiB
Plaintext
133 lines
7.8 KiB
Plaintext
|
|
(with_import "./match.kp"
|
|
(let (
|
|
; This is based on https://www.cs.cornell.edu/courses/cs3110/2020sp/a4/deletion.pdf
|
|
; and the figure references refer to it
|
|
; Insert is taken from the same paper, but is origional to Okasaki, I belive
|
|
|
|
; The tree has been modified slightly to take in a comparison function
|
|
; and override if insert replaces or not to allow use as a set or as a map
|
|
|
|
; I think this is actually pretty cool - instead of having a bunch of seperate ['B]
|
|
; be our leaf node, we use ['B] with all nils. This allows us to not use -B, as
|
|
; both leaf and non-leaf 'BB has the same structure with children! Also, we make
|
|
; sure to use empty itself so we don't make a ton of empties...
|
|
empty ['B nil nil nil]
|
|
E empty
|
|
EE ['BB nil nil nil]
|
|
|
|
size (rec-lambda recurse (t) (match t
|
|
~E 0
|
|
[c a x b] (+ 1 (recurse a) (recurse b))))
|
|
|
|
generic-foldl (rec-lambda recurse (f z t) (match t
|
|
~E z
|
|
[c a x b] (recurse f (f (recurse f z a) x) b)))
|
|
|
|
generic-contains? (rec-lambda recurse (t cmp v found not-found) (match t
|
|
~E (not-found)
|
|
[c a x b] (match (cmp v x) '< (recurse a cmp v found not-found)
|
|
'= (found x)
|
|
'> (recurse b cmp v found not-found))))
|
|
|
|
blacken (lambda (t) (match t
|
|
['R a x b] ['B a x b]
|
|
t t))
|
|
balance (lambda (t) (match t
|
|
; figures 1 and 2
|
|
['B ['R ['R a x b] y c] z d] ['R ['B a x b] y ['B c z d]]
|
|
['B ['R a x ['R b y c]] z d] ['R ['B a x b] y ['B c z d]]
|
|
['B a x ['R ['R b y c] z d]] ['R ['B a x b] y ['B c z d]]
|
|
['B a x ['R b y ['R c z d]]] ['R ['B a x b] y ['B c z d]]
|
|
; figure 8, double black cases
|
|
['BB ['R a x ['R b y c]] z d] ['B ['B a x b] y ['B c z d]]
|
|
['BB a x ['R ['R b y c] z d]] ['B ['B a x b] y ['B c z d]]
|
|
; already balenced
|
|
t t))
|
|
generic-insert (lambda (t cmp v replace) (let (
|
|
ins (rec-lambda ins (t) (match t
|
|
~E ['R t v t]
|
|
[c a x b] (match (cmp v x) '< (balance [c (ins a) x b])
|
|
'= (if replace [c a v b]
|
|
t)
|
|
'> (balance [c a x (ins b)]))))
|
|
) (blacken (ins t))))
|
|
|
|
rotate (lambda (t) (match t
|
|
; case 1, fig 6
|
|
['R ['BB a x b] y ['B c z d]] (balance ['B ['R ['B a x b] y c] z d])
|
|
['R ['B a x b] y ['BB c z d]] (balance ['B a x ['R b y ['B c z d]]])
|
|
; case 2, figure 7
|
|
['B ['BB a x b] y ['B c z d]] (balance ['BB ['R ['B a x b] y c] z d])
|
|
['B ['B a x b] y ['BB c z d]] (balance ['BB a x ['R b y ['B c z d]]])
|
|
; case 3, figure 9
|
|
['B ['BB a w b] x ['R ['B c y d] z e]] ['B (balance ['B ['R ['B a w b] x c] y d]) z e]
|
|
['B ['R a w ['B b x c]] y ['BB d z e]] ['B a w (balance ['B b x ['R c y ['B d z e]]])]
|
|
; fall through
|
|
t t))
|
|
|
|
redden (lambda (t) (match t
|
|
['B a x b] (if (and (= 'B (idx a 0)) (= 'B (idx b 0))) ['R a x b]
|
|
t)
|
|
t t))
|
|
|
|
min_delete (rec-lambda recurse (t) (match t
|
|
~E (error "min_delete empty tree")
|
|
['R ~E x ~E] [x E]
|
|
['B ~E x ~E] [x EE]
|
|
['B ~E x ['R a y b]] [x ['B a y b]]
|
|
[c a x b] (let ((v ap) (recurse a)) [v (rotate [c ap x b])])))
|
|
generic-delete (lambda (t cmp v) (let (
|
|
del (rec-lambda del (t v) (match t
|
|
; figure 3
|
|
~E t
|
|
; figure 4
|
|
['R ~E x ~E] (match (cmp v x) '= E
|
|
_ t)
|
|
['B ['R a x b] y ~E] (match (cmp v y) '< (rotate ['B (del ['R a x b] v) y ~E])
|
|
'= ['B a x b]
|
|
'> t)
|
|
; figure 5
|
|
['B ~E x ~E] (match (cmp v x) '= EE
|
|
_ t)
|
|
[c a x b] (match (cmp v x) '< (rotate [c (del a v) x b])
|
|
'= (let ([vp bp] (min_delete b))
|
|
(rotate [c a vp bp]))
|
|
'> (rotate [c a x (del b v)]))))
|
|
) (del (redden t) v)))
|
|
|
|
|
|
set-cmp (lambda (a b) (cond (< a b) '<
|
|
(= a b) '=
|
|
true '>))
|
|
set-empty empty
|
|
set-foldl generic-foldl
|
|
set-insert (lambda (t x) (generic-insert t set-cmp x false))
|
|
set-contains? (lambda (t x) (generic-contains? t set-cmp x (lambda (f) true) (lambda () false)))
|
|
set-remove (lambda (t x) (generic-delete t set-cmp x))
|
|
|
|
map-cmp (lambda (a b) (let (ak (idx a 0)
|
|
bk (idx b 0))
|
|
(cond (< ak bk) '<
|
|
(= ak bk) '=
|
|
true '>)))
|
|
map-empty empty
|
|
map-insert (lambda (t k v) (generic-insert t map-cmp [k v] true))
|
|
map-contains-key? (lambda (t k) (generic-contains? t map-cmp [k nil] (lambda (f) true) (lambda () false)))
|
|
map-get (lambda (t k) (generic-contains? t map-cmp [k nil] (lambda (f) (idx f 1)) (lambda () (error (str "didn't find key " k " in map " t)))))
|
|
map-get-or-default (lambda (t k d) (generic-contains? t map-cmp [k nil] (lambda (f) (idx f 1)) (lambda () d)))
|
|
map-get-with-default (lambda (t k d) (generic-contains? t map-cmp [k nil] (lambda (f) (idx f 1)) (lambda () (d))))
|
|
map-remove (lambda (t k) (generic-delete t map-cmp [k nil]))
|
|
|
|
; This could be 2x as efficent by being implmented on generic instead of map,
|
|
; as we wouldn't have to traverse once to find and once to insert
|
|
multimap-empty map-empty
|
|
multimap-insert (lambda (t k v) (map-insert t k (set-insert (map-get-or-default t k set-empty) v)))
|
|
multimap-get (lambda (t k) (map-get-or-default t k set-empty))
|
|
)
|
|
(provide set-empty set-foldl set-insert set-contains? set-remove
|
|
map-empty map-insert map-contains-key? map-get map-get-or-default map-get-with-default map-remove
|
|
multimap-empty multimap-insert multimap-get
|
|
size)
|
|
))
|