Purely Functional Red-Black Tree

ちょっと思い立って作ってみました。

最初はProgramming Praxisで示されている実装から出発したんですが、結果的にはHaskell版の実装に近い形に変化していきました。

アルゴリズム等の解説はWikipediaに譲ることにして、ここではマクロを活用して煩雑なパターンマッチをHaskellっぽく簡潔に行う方法など、テクニック的な話題や応用例を紹介したいと思います。

ツリー構文

まず、ツリーの構造を定義します。

(define-struct tree (color left node right))

(define empty (make-tree 'black #f (cons #f #f) #f))
(define (empty? t) (eq? t empty))

キーと値は cons で纏めて node として一体的に扱うことにします (最初は別々にしていたんですが、あまりにもコーディングが煩雑になったので…)。

構造体を利用する際に便利なのが match 構文です。

> (match (make-tree 'black empty '(1 . #f) empty)
    ((struct tree ('black (? empty?) x (? empty?)))
     "This is a black tree with both subtrees empty."))
"This is a black tree with both subtrees empty."

このように、アクセサ関数を使わずとも色の判定ができたりして、コードが簡潔になります。

でも、もっと簡潔にすることはできないでしょうか?

例えば

(match (T B E '(1 . #f) E)
  ((T B E x E)
   "This is a black tree with both subtrees empty."))

のように書ければ素晴らしいですよね。

define-match-expander というマクロを使うと、それが可能になります。

(define-match-expander T
  (lambda (stx)
    (syntax-case stx ()
      ((T C a x b)
       (syntax/loc stx
         (struct tree (C a x b))))))
  (lambda (stx)
    (syntax-case stx ()
      ((T C a x b)
       (syntax/loc stx
         (make-tree C a x b))))))

最初のラムダがマッチ構文内での展開形、2番目は式として呼ばれた場合の展開形の定義です。

(define-match-expander R
  (lambda (stx)
    (syntax-case stx ()
      (R (syntax 'red))))
  (lambda (stx)
    (syntax-case stx ()
      (R (syntax 'red)))))

(define-match-expander B
  (lambda (stx)
    (syntax-case stx ()
      (B (syntax 'black))))
  (lambda (stx)
    (syntax-case stx ()
      (B (syntax 'black)))))

空ツリー

(define-match-expander E
  (lambda (stx)
    (syntax-case stx ()
      (E (syntax/loc stx (? empty?)))))
  (lambda (stx)
    (syntax-case stx ()
      (E (syntax/loc stx empty)))))

R、B、E が、マクロでありながら値のように扱えることを確認してみましょう。

> (list R B E)
(red black #s(tree black #f (#f . #f) #f))

構文の準備ができると、例えば赤い木を黒くする関数が、このように書けるようになります。

(fun (blacken (and t (T B _ _ _))) t
  \| (blacken (T _ a x b))         (T B a x b))

パターンマッチとツリーの生成とが同じ形式になっていますね (2行目)。

1行目は、木が既に黒ければそのまま返すというパターンです。and が ML や Haskell の as パターンと同じ意味で使えるということです。

もちろんネスティングも可能です。

> (match (T R E '(1) (T B (T R E '(2) E) '(3) E))
    ((T R _ x (T B (T R _ y _) z _))
     (append x y z)))
(1 2 3)

このHaskell風の構文糖衣が無ければいかに煩雑なコードになるか、想像に難くないでしょう。

キーの比較

ツリーのバランスを保ったり、要素を検索したりする上で必要なのが、キーの比較関数です。

数値や文字列、シンボルなど、比較可能なものであれば何でもキーにできるようにしたいんですが、比較方法を利用側で指定するのはやや面倒なんですね。

そこで、型変換ライブラリと同じような要領で、予め既知のデータ型については比較関数を用意しておいて、キーが与えられた時にそれに合ったものを選ぶ、という方式を考えました。

(fun (hany f h)
  (prompt
    (hash-for-each h
                   (fn (k v)
                     (aif (f k v)
                       (control _ it))))
    #f))

(fun (cmp = < x y)
  (if (= x y) '=
      (< x y) '<
              '>))

(define-values (install-cmp comparator)
  (let ((h (make-hasheq)))
    (values (fn (? = <)
              (hash-set! h ? (cmp = <)))
            (fn (x)
              (hany (fn (? cmp)
                      (and (? x) cmp))
                    h)))))

データ型に応じた比較関数を以下のようにインストールしておきます。

(install-cmp number? = <)
(install-cmp symbol? eq? (fn (x y)
                           (string<? (symbol->string x)
                                     (symbol->string y))))
...

テスト

> (let ((cmp ((comparator 0) 0)))
    (values (cmp 1) (cmp 0) (cmp -1)))
<
=
>

これにより、ユーザーに対して具体的な比較方法を意識させない仕方でツリー生成関数を提供できることになります。すなわち、比較関数ではなく代表的なキーの値を与えさせるわけです。

(fun (make-tree~ k) (cons (comparator k) empty))

(make-tree は構造体の生成関数なので、ライブラリ中ではチルダを付けておいて、エクスポート時にチルダを外すようにします)

ツリーの利用者はこのインターフェース関数を呼び出して空のツリーと比較関数を (無意識に) 作ります。

なおライブラリ側では、全てのツリー操作を通じて最初に作られた比較関数を使い回すようにすべきでしょう。

> (let* ((l '(4 1 2 5 3))
         (t (make-tree (car l))))
    (eq? (car (foldl (lambda (x t)
                       (tree-set t x '()))
                     t
                     l))
         (car t)))
#t

ツリー更新

ツリーに新しいノードを追加するための基本関数を、このように定義してみました。

(fun (tree-set~ reduce (cons cmp t) k v)
  (cons cmp
        (blacken
         (let loop (((and t
                          (T C a (and x (cons k~ v~)) b))
                     t))
           (if (empty? t)
               (T R E (cons k v) E)
               (case (cmp k k~)
                 (< (if (black? t)
                        (balance (loop a) x b)
                        (T R (loop a) x b)))
                 (= (T C a (cons k (reduce k v~ v)) b))
                 (> (if (black? t)
                        (balance a x (loop b))
                        (T R a x (loop b))))))))))

既に同じキーの要素が存在した場合の対応をパラメータ化することで (reduce)、以下のように自在に派生関数を作り出すことができます。

上書き版

(define tree-set (tree-set~ (fn (k old new) new)))

更新版

(fun (tree-add + t k v)
  (tree-set~ (fn (k old new) (+ old new)) t k v))

追加版

(fun (tree-cons t k v)
  (tree-add append t k (list v)))

なお balance 関数は、Haskell 版そのまんまです。

(define-syntax defmatch
  (syntax-rules ()
    ((defmatch name ((pat ...) . body) ...)
     (define name
       (match-lambda** ((pat ...) . body) ...)))))

(defmatch balance
  (((T R a x b) y (T R c z d))
   (T R (T B a x b) y (T B c z d)))
  (((T R (T R a x b) y c) z d)
   (T R (T B a x b) y (T B c z d)))
  (((T R a x (T R b y c)) z d)
   (T R (T B a x b) y (T B c z d)))
  ((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 (T R (T R b y c) z d))
   (T R (T B a x b) y (T B c z d)))
  ((a x b)
   (T B a x b)))

応用例 (素因数分解)

tree-add を利用する例として、同じオブジェクトの出現回数を数える、というケースが考えられます。ここでは整数を素因数のリストに変換して、それをさらにSchemeの数式に変換する、という問題を解きます。

(fun (factorize~ n)
  (let loop ((n n) (d 2) (l '()))
    (if (< (/ n d) d)
        (cons n l)
        (receive (q r) (quotient/remainder n d)
          (if (zero? r)
              (loop q d (cons d l))
              (loop n (+ (if (= d 2) 1 2) d) l))))))

中身はあまり重要ではないので結果だけ見てください。

> (factorize~ 1984)
(31 2 2 2 2 2 2)

これにシンボル * を cons すればそのまま掛け算の式になりますが、そうではなく、同じ数の掛け算は指数の形に変換したいです。

こんな時、tree-add が使えます。

(fun (factorize n)
  (let ((l (map (fn ((cons n p))
                  (if (= p 1) n `(expt ,n ,p)))
                (tree->alist
                 (foldl (fn (x t)
                          (tree-add + (or t (make-tree x)) x 1))
                        #f
                        (factorize~ n))))))
    (if (null? (cdr l))
        (car l)        ; prime number
        (cons '* l))))

いちおう検算

> (apply values
         (map ((fn (x f) (f x)) 1984)
              (list factorize~
                    factorize
                    (compose eval factorize))))
(31 2 2 2 2 2 2)
(* (expt 2 6) 31)
1984

2番目の値がツリーの活用の成果です。細かいですが、2の指数の項が先に来ているのは偶然ではなく、キーの並び順通りにリスト化された結果です。

なお、ツリーからリストへの変換は以下のように定義しています。

(fun (tree-fold (cons _ t) f seed)
  (let loop (((and t
                   (T _ a (cons k v) b))
              t)
             (seed seed))
    (if (empty? t)
        seed
        (loop a
              (f k v (loop b seed))))))

(fun (tree-map f t)
  (tree-fold t (fn (k v l) (cons (f k v) l)) '()))

(define tree-keys (tree-map (fn (k _) k)))
(define tree-values (tree-map (fn (_ v) v)))
(define tree->alist (tree-map cons))

既存ノードの更新

キーが既に存在する場合のみ値を更新したい、という要求が生じることもあります。

これを単純に、ツリーを検索して見つかれば tree-set で更新、というやり方にしてしまうと、同じ検索を2回繰り返すという無駄が生じます。

また、見つからなければ何もしない、というのは tree-set~ では表現できないことなので、別な関数が必要になります。

(fun (tree-reset new (cons cmp t) k)
  (cons cmp
        (or (let loop (((and t
                             (T C a (and x (cons k~ v)) b))
                        t))
              (if (empty? t)
                  #f
                  (case (cmp k k~)
                    (< (aif (loop a) (T C it x b)))
                    (= (T C a (cons k (new v)) b))
                    (> (aif (loop b) (T C a x it))))))
            t)))

ツリーが空になった (キーが見つからなかった) 場合に、#f の代わりに空ツリーを返すようにすれば or や aif は不要になるんですが、不要なツリー生成をしないためにこのように工夫してあります。



マクロ: anaphora.ss arcfun.ss arcif.ss letfun.ss mlfun.ss