Skip to content

Instantly share code, notes, and snippets.

@loganmhb
Created Oct 2, 2019
Embed
What would you like to do?
Persistent sorted set implemented as a red-black tree in Common Lisp
;; todo: how do you actually manage dependencies?
(ql:quickload "generic-cl")
(ql:quickload "trivia")
;; Persistent red-black tree implemnetation per Okasaki's paper:
;; https://www.cs.tufts.edu/~nr/cs257/archive/chris-okasaki/redblack99.pdf
(defpackage sset
(:use generic-cl))
(in-package sset)
(defstruct rbtree color left element right)
(defstruct empty-rbtree)
(defun empty-set () (make-empty-rbtree))
(defun set-member (rbtree item)
(trivia:match rbtree
((empty-rbtree) nil)
((rbtree (color 'red)
(left l)
(element e)
(right r))
(cond
((< item e) (set-member l item))
((equalp item e) t)
((> item e) (set-member r item))))))
(defun balance (rbtree)
(trivia:match rbtree
;; match all black subtrees with red child with another red child, so
;; we can reconstruct them as a red subtree with two black
;; children (Okasaki fig. 1)
((or (rbtree :color 'black
:left (rbtree :color 'red
:left (rbtree :color 'red
:left a
:element x
:right b)
:element y
:right c)
:element z
:right d)
(rbtree :color 'black
:left (rbtree :color 'red
:left a
:element x
:right (rbtree :color 'red
:left b
:element y
:right c))
:element z
:right d)
(rbtree :color 'black
:left a
:element x
:right (rbtree :color 'red
:left (rbtree :color 'red
:left b
:element y
:right c)
:element z
:right d))
(rbtree :color 'black
:left a
:element x
:right (rbtree :color 'red
:left b
:element y
:right (rbtree :color 'red
:left c
:element z
:right d))))
(make-rbtree :color 'red
:left (make-rbtree :color 'black
:left a
:element x
:right b)
:element y
:right (make-rbtree :color 'black
:left c
:element z
:right d)))
(rbtree rbtree)))
(defun add (rbtree item)
(labels ((ins (tree)
(trivia:match tree
((empty-rbtree) (make-rbtree :color 'red
:left (make-empty-rbtree)
:element item
:right (make-empty-rbtree)))
((rbtree :color color :left left :element element :right right)
(cond
((< item element) (balance (make-rbtree :color color
:left (ins left)
:element element
:right right)))
;; return a new tree even though it's identical,
;; because we will be changing the color to black
((= item element) (make-rbtree :color color
:left left
:element element
:right right))
((> item element) (balance (make-rbtree :color color
:left left
:element element
:right (ins right)))))))))
(let ((new-tree (ins rbtree)))
(setf (slot-value new-tree 'color) 'black)
new-tree)))
(defun elements (rbtree)
(trivia:match rbtree
((empty-rbtree) '())
((rbtree :color _ :left left :element elt :right right)
(concatenate (elements left) (list elt) (elements right)))))
(defun list->rbtree (list)
(reduce #'add list :initial-value (make-empty-rbtree)))
(defun range (rbtree &key from to)
;; todo: iterator, not list?
(trivia:match rbtree
((empty-rbtree) '())
((rbtree :left l :element x :right r)
(concatenate (when (< from x)
(range l :from from :to to))
(when (and (<= from x) (< x to))
(list x))
(when (< x to)
(range r :from from :to to))))))
;; (range (list->rbtree '(1 2 3 4 5 6 7 8 9 10)) :from 4 :to 7)
;; (elements (list->rbtree '("bear" "apple" "pear" "orange")))
;; trees should be sorted:
;; (time
;; (= (elements (list->rbtree (loop for x from 10000 downto 0 collect x)))
;; (elements (list->rbtree (loop for x from 0 to 10000 collect x)))))
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment