479 lines
12 KiB
Racket
479 lines
12 KiB
Racket
#lang racket/base
|
|
;; Matt Might's red-black tree code from
|
|
;; http://matt.might.net/articles/red-black-delete/code/sorted-map.rkt
|
|
;;
|
|
;; Modified by Tony Garnock-Jones, July 2014:
|
|
;; - trees are hashconsed
|
|
;; - sorted-map-size is made constant-time
|
|
|
|
(provide (struct-out sorted-map)
|
|
sorted-map-empty
|
|
sorted-map-modify-at
|
|
sorted-map-insert
|
|
sorted-map-insert*
|
|
sorted-map-to-tree
|
|
sorted-map-to-alist
|
|
sorted-map-submap?
|
|
sorted-map-get
|
|
sorted-map-size
|
|
sorted-map-max
|
|
sorted-map-delete
|
|
sorted-map-has-key?
|
|
sorted-map-keys
|
|
sorted-map-values
|
|
)
|
|
|
|
(require "canonicalize.rkt")
|
|
(require "memoize.rkt")
|
|
(require racket/set)
|
|
|
|
; A purely functional sorted-map library.
|
|
|
|
; Provides logarithmic insert, update, get & delete.
|
|
|
|
; Based on Okasaki's red-black trees
|
|
; with purely functional red-black delete.
|
|
|
|
; Author: Matthew Might
|
|
; Site: http://matt.might.net/
|
|
; Page: http://matt.might.net/articles/red-black-delete/
|
|
|
|
(require (except-in racket/match define/match))
|
|
|
|
; Syntactic sugar for define forms
|
|
; with match as their body:
|
|
(define-syntax define/match
|
|
(syntax-rules ()
|
|
[(_ (id name) clause ...)
|
|
; =>
|
|
(define (id name)
|
|
(match name clause ...))]))
|
|
|
|
(define-syntax define/match*
|
|
(syntax-rules ()
|
|
[(_ (id name ...) clause ...)
|
|
; =>
|
|
(define (id name ...)
|
|
(match* (name ...)
|
|
clause ...))]))
|
|
|
|
; A form for matching the result of a comparison:
|
|
;; tonyg 20140718: changed to use the order? convention from data/order
|
|
(define-syntax switch-compare
|
|
(syntax-rules (= < >)
|
|
[(_ (cmp v1 v2)
|
|
[< action1 ...]
|
|
[= action2 ...]
|
|
[> action3 ...])
|
|
; =>
|
|
(let ((dir (cmp v1 v2)))
|
|
(case dir
|
|
[(<) action1 ...]
|
|
[(=) action2 ...]
|
|
[(>) action3 ...]))]))
|
|
|
|
;; tonyg 20140718: for hash-consing, we have to be able to compare
|
|
;; trees using equal?, which necessitates use of #:transparent in our
|
|
;; struct definitions.
|
|
|
|
; Struct definition for sorted-map:
|
|
(define-struct sorted-map (compare) #:transparent)
|
|
|
|
; Internal nodes:
|
|
(define-struct (T sorted-map)
|
|
(color left key value right) #:transparent)
|
|
|
|
; Leaf nodes:
|
|
(define-struct (L sorted-map) () #:transparent)
|
|
|
|
; Double-black leaf nodes:
|
|
(define-struct (BBL sorted-map) () #:transparent)
|
|
|
|
|
|
; Color manipulators.
|
|
|
|
; Turns a node black.
|
|
(define/match (blacken node)
|
|
[(T cmp _ l k v r) (canonicalize (T cmp 'B l k v r))]
|
|
[(BBL cmp) (canonicalize (L cmp))]
|
|
[(L _) node])
|
|
|
|
; Turns a node red.
|
|
(define/match (redden node)
|
|
[(T cmp _ l k v r) (canonicalize (T cmp 'R l k v r))]
|
|
[(L _) (error "Can't redden leaf.")])
|
|
|
|
|
|
; Color arithmetic.
|
|
(define/match (black+1 color-or-node)
|
|
[(T cmp c l k v r) (canonicalize (T cmp (black+1 c) l k v r))]
|
|
[(L cmp) (canonicalize (BBL cmp))]
|
|
['-B 'R]
|
|
['R 'B]
|
|
['B 'BB])
|
|
|
|
(define/match (black-1 color-or-node)
|
|
[(T cmp c l k v r) (canonicalize (T cmp (black-1 c) l k v r))]
|
|
[(BBL cmp) (canonicalize (L cmp))]
|
|
['R '-B]
|
|
['B 'R]
|
|
['BB 'B])
|
|
|
|
|
|
|
|
; Creates an empty map:
|
|
(define (sorted-map-empty compare)
|
|
(canonicalize (L compare)))
|
|
|
|
|
|
;; Custom patterns.
|
|
|
|
; Matches internal nodes:
|
|
(define-match-expander T!
|
|
(syntax-rules ()
|
|
[(_) (T _ _ _ _ _ _)]
|
|
[(_ l r) (T _ _ l _ _ r)]
|
|
[(_ c l r) (T _ c l _ _ r)]
|
|
[(_ l k v r) (T _ _ l k v r)]
|
|
[(_ c l k v r) (T _ c l k v r)]))
|
|
|
|
; Matches leaf nodes:
|
|
(define-match-expander L!
|
|
(syntax-rules ()
|
|
[(_) (L _)]))
|
|
|
|
; Matches black nodes (leaf or internal):
|
|
(define-match-expander B
|
|
(syntax-rules ()
|
|
[(_) (or (T _ 'B _ _ _ _)
|
|
(L _))]
|
|
[(_ cmp) (or (T cmp 'B _ _ _ _)
|
|
(L cmp))]
|
|
[(_ l r) (T _ 'B l _ _ r)]
|
|
[(_ l k v r) (T _ 'B l k v r)]
|
|
[(_ cmp l k v r) (T cmp 'B l k v r)]))
|
|
|
|
; Matches red nodes:
|
|
(define-match-expander R
|
|
(syntax-rules ()
|
|
[(_) (T _ 'R _ _ _ _)]
|
|
[(_ cmp) (T cmp 'R _ _ _ _)]
|
|
[(_ l r) (T _ 'R l _ _ r)]
|
|
[(_ l k v r) (T _ 'R l k v r)]
|
|
[(_ cmp l k v r) (T cmp 'R l k v r)]))
|
|
|
|
; Matches negative black nodes:
|
|
(define-match-expander -B
|
|
(syntax-rules ()
|
|
[(_) (T _ '-B _ _ _ _)]
|
|
[(_ cmp) (T cmp '-B _ _ _ _)]
|
|
[(_ l k v r) (T _ '-B l k v r)]
|
|
[(_ cmp l k v r) (T cmp '-B l k v r)]))
|
|
|
|
; Matches double-black nodes (leaf or internal):
|
|
(define-match-expander BB
|
|
(syntax-rules ()
|
|
[(_) (or (T _ 'BB _ _ _ _)
|
|
(BBL _))]
|
|
[(_ cmp) (or (T cmp 'BB _ _ _ _)
|
|
(BBL _))]
|
|
[(_ l k v r) (T _ 'BB l k v r)]
|
|
[(_ cmp l k v r) (T cmp 'BB l k v r)]))
|
|
|
|
(define/match (double-black? node)
|
|
[(BB) #t]
|
|
[_ #f])
|
|
|
|
|
|
; Turns a black-balanced tree with invalid colors
|
|
; into a black-balanced tree with valid colors:
|
|
(define (balance-node node)
|
|
(define cmp (sorted-map-compare node))
|
|
(match node
|
|
[(or (T! (or 'B 'BB) (R (R a xk xv b) yk yv c) zk zv d)
|
|
(T! (or 'B 'BB) (R a xk xv (R b yk yv c)) zk zv d)
|
|
(T! (or 'B 'BB) a xk xv (R (R b yk yv c) zk zv d))
|
|
(T! (or 'B 'BB) a xk xv (R b yk yv (R c zk zv d))))
|
|
; =>
|
|
(canonicalize (T cmp
|
|
(black-1 (T-color node))
|
|
(canonicalize (T cmp 'B a xk xv b))
|
|
yk
|
|
yv
|
|
(canonicalize (T cmp 'B c zk zv d))))]
|
|
|
|
[(BB a xk xv (-B (B b yk yv c) zk zv (and d (B))))
|
|
; =>
|
|
(canonicalize (T cmp
|
|
'B
|
|
(canonicalize (T cmp 'B a xk xv b))
|
|
yk
|
|
yv
|
|
(balance cmp 'B c zk zv (redden d))))]
|
|
|
|
[(BB (-B (and a (B)) xk xv (B b yk yv c)) zk zv d)
|
|
; =>
|
|
(canonicalize (T cmp
|
|
'B
|
|
(balance cmp 'B (redden a) xk xv b)
|
|
yk
|
|
yv
|
|
(canonicalize (T cmp 'B c zk zv d))))]
|
|
|
|
[else node]))
|
|
|
|
(define (balance cmp c l k v r)
|
|
(balance-node (canonicalize (T cmp c l k v r))))
|
|
|
|
|
|
; Moves to a location in the map and
|
|
; peformes an update with the function:
|
|
;; tonyg 20140718 added on-missing argument
|
|
(define (sorted-map-modify-at node key f [on-missing (lambda () #f)])
|
|
|
|
(define (internal-modify-at node key f)
|
|
(match node
|
|
[(T cmp c l k v r)
|
|
; =>
|
|
(switch-compare (cmp key k)
|
|
[< (balance cmp c (internal-modify-at l key f) k v r)]
|
|
[= (canonicalize (T cmp c l k (f k v) r))]
|
|
[> (balance cmp c l k v (internal-modify-at r key f))])]
|
|
|
|
[(L cmp)
|
|
; =>
|
|
(canonicalize (T cmp 'R node key (f key (on-missing)) node))]))
|
|
|
|
(blacken (internal-modify-at node key f)))
|
|
|
|
|
|
; Inserts an element into the map
|
|
(define (sorted-map-insert node key value)
|
|
(sorted-map-modify-at node key (lambda (k v) value)))
|
|
|
|
|
|
; Inserts several elements into the map:
|
|
(define (sorted-map-insert* node keys values)
|
|
(if (or (not (pair? keys))
|
|
(not (pair? values)))
|
|
node
|
|
(sorted-map-insert*
|
|
(sorted-map-insert node (car keys) (car values))
|
|
(cdr keys) (cdr values))))
|
|
|
|
|
|
; Coverts a sorted map into a tree:
|
|
(define/match (sorted-map-to-tree node)
|
|
[(L!) 'L]
|
|
[(T! c l k v r) `(,c ,(sorted-map-to-tree l) ,k ,v ,(sorted-map-to-tree r))]
|
|
[else node])
|
|
|
|
|
|
; Converts a sorted map into an alist:
|
|
(define (sorted-map-to-alist node)
|
|
|
|
(define (sorted-map-prepend-as-alist node alist)
|
|
(match node
|
|
[(T! l k v r)
|
|
; =>
|
|
(sorted-map-prepend-as-alist
|
|
l
|
|
(cons (cons k v)
|
|
(sorted-map-prepend-as-alist r alist)))]
|
|
|
|
[(L _)
|
|
; =>
|
|
alist]))
|
|
|
|
(sorted-map-prepend-as-alist node '()))
|
|
|
|
|
|
; Tests whether this map is a submap of another map:
|
|
(define (sorted-map-submap? map1 map2 #:by [by equal?])
|
|
(define amap1 (sorted-map-to-alist map1))
|
|
(define amap2 (sorted-map-to-alist map2))
|
|
(define cmp (sorted-map-compare map1))
|
|
|
|
(define (compare-alists amap1 amap2)
|
|
(match* (amap1 amap2)
|
|
[(`((,k1 . ,v1) . ,rest1)
|
|
`((,k2 . ,v2) . ,rest2))
|
|
; =>
|
|
(switch-compare (cmp k1 k2)
|
|
[< #f]
|
|
[= (and (by v1 v2) (compare-alists rest1 rest2))]
|
|
[> (compare-alists amap1 rest2)])]
|
|
|
|
[('() '()) #t]
|
|
|
|
[(_ '()) #f]
|
|
|
|
[('() _) #t]))
|
|
|
|
(compare-alists amap1 amap2))
|
|
|
|
|
|
; Gets an element from a sorted map:
|
|
;; tonyg 20140718 add on-missing argument
|
|
(define (sorted-map-get node key [on-missing (lambda () #f)])
|
|
(let walk ((node node))
|
|
(match node
|
|
[(L!) (on-missing)]
|
|
|
|
[(T cmp c l k v r)
|
|
; =>
|
|
(switch-compare (cmp key k)
|
|
[< (walk l)]
|
|
[= v]
|
|
[> (walk r)])])))
|
|
|
|
|
|
; Returns the size of the sorted map:
|
|
;; tonyg 20140718 this is memoized to run in O(1) for every smap
|
|
(define sorted-map-size
|
|
(memoize1
|
|
(lambda (smap)
|
|
(match smap
|
|
[(T! l r) (+ 1 (sorted-map-size l)
|
|
(sorted-map-size r))]
|
|
[(L!) 0]))))
|
|
|
|
|
|
; Returns the maxium (key . value) pair:
|
|
(define/match (sorted-map-max node)
|
|
[(T! _ k v (L!)) (cons k v)]
|
|
[(T! _ r) (sorted-map-max r)])
|
|
|
|
|
|
; Performs a check to see if both invariants are met:
|
|
(define (sorted-map-is-legal? node)
|
|
|
|
; Calculates the max black nodes on path:
|
|
(define/match (max-black-height node)
|
|
[(T! c l r)
|
|
; =>
|
|
(+ (if (eq? c 'B) 1 0) (max (max-black-height l)
|
|
(max-black-height r)))]
|
|
|
|
[(L!) 1])
|
|
|
|
; Calculates the min black nodes on a path:
|
|
(define/match (min-black-height node)
|
|
[(T! c l r)
|
|
; =>
|
|
(+ (if (eq? c 'B) 1 0) (min (min-black-height l)
|
|
(min-black-height r)))]
|
|
|
|
[(L!) 1])
|
|
|
|
; Is this tree black-balanced?
|
|
(define (black-balanced? node)
|
|
(= (max-black-height node)
|
|
(min-black-height node)))
|
|
|
|
; Does this tree contain a red child of red?
|
|
(define/match (no-red-red? node)
|
|
[(or (B l r)
|
|
(R (and l (B)) (and r (B))))
|
|
; =>
|
|
(and (no-red-red? l) (no-red-red? r))]
|
|
|
|
[(L!) #t]
|
|
[else #f])
|
|
|
|
(let ((colored? (no-red-red? node))
|
|
(balanced? (black-balanced? node)))
|
|
(and colored? balanced?)))
|
|
|
|
|
|
; Deletes a key from this map:
|
|
(define (sorted-map-delete node key)
|
|
|
|
(define cmp (sorted-map-compare node))
|
|
|
|
; Finds the node to be removed:
|
|
(define/match (del node)
|
|
[(T! c l k v r)
|
|
; =>
|
|
(switch-compare (cmp key k)
|
|
[< (bubble c (del l) k v r)]
|
|
[= (remove node)]
|
|
[> (bubble c l k v (del r))])]
|
|
|
|
[else node])
|
|
|
|
; Removes this node; it might
|
|
; leave behind a double-black node:
|
|
(define/match (remove node)
|
|
; Leaves are easiest to kill:
|
|
[(R (L!) (L!)) (canonicalize (L cmp))]
|
|
[(B (L!) (L!)) (canonicalize (BBL cmp))]
|
|
|
|
; Killing a node with one child;
|
|
; parent or child is red:
|
|
[(or (R child (L!))
|
|
(R (L!) child))
|
|
; =>
|
|
child]
|
|
|
|
[(or (B (R l k v r) (L!))
|
|
(B (L!) (R l k v r)))
|
|
; =>
|
|
(canonicalize (T cmp 'B l k v r))]
|
|
|
|
; Killing a black node with one black child:
|
|
[(or (B (L!) (and child (B)))
|
|
(B (and child (B)) (L!)))
|
|
; =>
|
|
(black+1 child)]
|
|
|
|
; Killing a node with two sub-trees:
|
|
[(T! c (and l (T!)) (and r (T!)))
|
|
; =>
|
|
(match-let (((cons k v) (sorted-map-max l))
|
|
(l* (remove-max l)))
|
|
(bubble c l* k v r))])
|
|
|
|
; Kills a double-black, or moves it to the top:
|
|
(define (bubble c l k v r)
|
|
(cond
|
|
[(or (double-black? l) (double-black? r))
|
|
; =>
|
|
(balance cmp (black+1 c) (black-1 l) k v (black-1 r))]
|
|
|
|
[else (canonicalize (T cmp c l k v r))]))
|
|
|
|
; Removes the max node:
|
|
(define/match (remove-max node)
|
|
[(T! l (L!)) (remove node)]
|
|
[(T! c l k v r ) (bubble c l k v (remove-max r))])
|
|
|
|
; Delete the key, and color the new root black:
|
|
(blacken (del node)))
|
|
|
|
|
|
;; tonyg 20140718 True iff key is in node
|
|
(define (sorted-map-has-key? node key)
|
|
(let walk ((node node))
|
|
(match node
|
|
[(L!) #f]
|
|
[(T cmp c l k v r)
|
|
(switch-compare (cmp key k)
|
|
[< (walk l)]
|
|
[= #t]
|
|
[> (walk r)])])))
|
|
|
|
;; tonyg 20140718 Retrieve a set of the keys of smap
|
|
(define (sorted-map-keys smap [empty-set (set)])
|
|
(let walk ((node smap) (acc empty-set))
|
|
(match node
|
|
[(T! l k v r) (walk l (set-add (walk r acc) k))]
|
|
[(L _) acc])))
|
|
|
|
;; tonyg 20140718 Retrieve a list of the values of smap
|
|
(define (sorted-map-values smap)
|
|
(let walk ((node smap) (acc '()))
|
|
(match node
|
|
[(T! l k v r) (walk l (cons v (walk r acc)))]
|
|
[(L _) acc])))
|