Common Lispで赤黒木

Haskellでの赤黒木の実装の美しさに感動した

data Color = R | B
data RedBlackSet a = E | T Color (RedBlackSet a) a (RedBlackSet a)

balance B (T R (T R a x b) y c) z d = T R (T B a x b) y (T B c z d)
balance B (T R a x (T R b y c)) z d = T R (T B a x b) y (T B c z d)
balance B a x (T R (T R b y c) z d) = T R (T B a x b) y (T B c z d)
balance B a x (T R b y (T R c z d)) = T R (T B a x b) y (T B c z d)
balance color a x b = T color a x b

insert x s = T B a y b
  where 
    T _ a y b = ins s
    ins E = T R E x E
    ins s@(T color a y b) 
      | x < y  = balance color (ins a) y b
      | x > y  = balance color a y (ins b)
      | True   = s

(http://www.sampou.org/cgi-bin/haskell.cgi?Programming%3a%b6%cc%bc%ea%c8%a2%3a%a4%bd%a4%ce%c2%be#H-179mxttより)
この美しさは、

T Color a x b

としたとき

aに含まれる要素 < x < bに含まれる要素

が成り立つため、「赤の子は黒」条件が破れた木を赤黒木に戻す関数「balance」が

balance B (T R (T R a x b) y c) z d = T R (T B a x b) y (T B c z d)
balance B (T R a x (T R b y c)) z d = T R (T B a x b) y (T B c z d)
balance B a x (T R (T R b y c) z d) = T R (T B a x b) y (T B c z d)
balance B a x (T R b y (T R c z d)) = T R (T B a x b) y (T B c z d)

と「a x b y c z d」を左から右に順に書き並べれば定義できることから来ると思う。

CommonLispでもパターンマッチさえ実装すればほぼ同様に書けて

(defun balance (tree)
  (pattern-match tree
    ((:pattern (B (R (R i x j) y k) z l) :variable (x y z i j k l))
     `(R (B ,i ,x ,j) ,y (B ,k ,z ,l)))
    ((:pattern (B (R i x (R j y k)) z l) :variable (x y z i j k l))
     `(R (B ,i ,x ,j) ,y (B ,k ,z ,l)))
    ((:pattern (B i x (R (R j y k) z l)) :variable (x y z i j k l))
     `(R (B ,i ,x ,j) ,y (B ,k ,z ,l)))   
    ((:pattern (B i x (R j y (R k z l))) :variable (x y z i j k l))
     `(R (B ,i ,x ,j) ,y (B ,k ,z ,l)))
    (:otherwise tree)))
                 
(defun rb-insert% (tree z cmp)
  (pattern-match tree
    ((:pattern (color right element left) :variable (color right element left))
     (balance (if (funcall cmp z element)
                  `(,color ,right ,element ,(rb-insert% left z cmp))
                  `(,color ,(rb-insert% right z cmp) ,element ,left))))
    ((:pattern nil)
     `(R nil ,z nil))))

(defun rb-insert (tree z cmp)
  `(B ,@(rest (rb-insert% tree z cmp))))

となる。(Red Black Tree for Common Lisp. · GitHub)

CL-USER 18 > nil
NIL
CL-USER 19 > (rb-insert * (random 100)  #'<)
(B NIL 5 NIL)
CL-USER 20 > (rb-insert * (random 100)  #'<) ;;以下省略
(B (R NIL 32 NIL) 5 NIL)
(B (B NIL 51 NIL) 32 (B NIL 5 NIL))
(B (B NIL 51 NIL) 32 (B (R NIL 23 NIL) 5 NIL))
(B (B (R NIL 63 NIL) 51 NIL) 32 (B (R NIL 23 NIL) 5 NIL))
(B (R (B NIL 63 NIL) 63 (B NIL 51 NIL)) 32 (B (R NIL 23 NIL) 5 NIL))
(B (R (B (R NIL 84 NIL) 63 NIL) 63 (B NIL 51 NIL)) 32 (B (R NIL 23 NIL) 5 NIL))
(B (R (B (R NIL 84 NIL) 63 NIL) 63 (B NIL 51 (R NIL 48 NIL))) 32 (B (R NIL 23 NIL) 5 NIL))
(B (R (B (R NIL 84 NIL) 63 NIL) 63 (B NIL 51 (R NIL 48 NIL))) 32 (R (B NIL 23 NIL) 16 (B NIL 5 NIL)))

と正しく挿入できている。

パターンマッチ用のマクロは以下

(defmacro with-gensyms (gensym-variables &body body)
  `(let ,(loop for gensym-variable in gensym-variables
               collect `(,gensym-variable (gensym (symbol-name ',gensym-variable))))
     ,@body))

(defun matcher (target variables values pattern)
  (cond ((consp pattern)
         (with-gensyms (cartarget cdrtarget)
           `(and (consp ,target)
                 (let ((,cartarget (car ,target))
                       (,cdrtarget (cdr ,target)))
                   (and ,(matcher cartarget variables values (car pattern))
                        ,(matcher cdrtarget variables values (cdr pattern)))))))
        ((symbolp pattern)
         (cond ((member pattern variables)
                `(progn (setf ,pattern ,target) t))
               ((member pattern values)
                `(equal ,target ,pattern))
               (:otherwise
                `(eq ,target ',pattern))))
        (:otherwise
         `(equal ,target ,pattern))))

(defmacro if-matches (target-expr (&key (variable nil) (value nil) (guard t) pattern) then else)
  (with-gensyms (target)
    `(let ,variable
       (let ((,target ,target-expr))
         (if (and ,(matcher target variable value pattern)
                  ,guard)
             ,then
             ,else)))))

(defmacro pattern-match (target-expr &body match-clauses)
  (with-gensyms (target block-name)
    `(block ,block-name
       (let ((,target ,target-expr))
         ,@(loop for match-clause in match-clauses
                 collect (if (eq (first match-clause) :otherwise)
                             `(return-from ,block-name ,(second match-clause))
                             `(if-matches ,target ,(first match-clause)
                                      (return-from ,block-name ,(second match-clause))
                                      nil)))))))