tree-seq と flatten - 水底で思うこと

◆tree-seq をちゃんと理解したい

理解していなかったのでいろいろやってみた。

(defn tree-seq
  "Returns a lazy sequence of the nodes in a tree, via a depth-first walk.
   branch? must be a fn of one arg that returns true if passed a node
   that can have children (but may not).  children must be a fn of one
   arg that returns a sequence of the children. Will only be called on
   nodes for which branch? returns true. Root is the root node of the
  tree."
  {:added "1.0"}
  [branch? children root]
  (let [walk (fn walk [node]
               (lazy-seq
                (cons node
                      (when (branch? node)
                        (mapcat walk (children node))))))]
    (walk root)))

引数はそれぞれ、

root: tree構造のトップレベルノード
branch?: ノードを引数とし、そのノードがleaf(末端)でないとき true に評価される関数
children: ノードを引数とし、そのノードの子ノードをシーケンスにして返す関数

となります。

◆ノードを record で定義して、木を作る。

木を対象とするようなので木を作る。

(defrecord Node [value left right])

(defn node
  "ノード生成ヘルパー"
  [& kvs]
  (apply assoc (Node. nil nil nil) kvs))
 
(def tree
     (node :value "root"
           :left (node :value "L"
                       :left (node :value "LL")
                       :right (node :value "LR"))
           :right (node :value "R"
                        :left (node :value "RL")
                        :right (node :value "RR"))))

tree-seq に適用するには、あと二つ関数が必要なのでそれも作る

(defn branch?
  "左右の子ノードの少なくとも一つが nil じゃない"
  [node]
  (or (:left node) (:right node)))

(defn children
  "nil じゃない子ノードをシーケンスにする"
  [node]
  (remove nil? [(:left node) (:right node)]))

これを tree-seq に適用する

(tree-seq branch? children tree)

評価値は次のような「部分木を要素とするシーケンス」になります。みやすいように要素毎に空行いれてます。

({:value "root",
  :left {:value "L",
         :left {:value "LL", :left nil, :right nil},
         :right {:value "LR", :left nil, :right nil}},
  :right {:value "R",
          :left {:value "RL", :left nil, :right nil},
          :right {:value "RR", :left nil, :right nil}}}
 
 {:value "L",
  :left {:value "LL", :left nil, :right nil},
  :right {:value "LR", :left nil, :right nil}}
 
 {:value "LL", :left nil, :right nil}
 
 {:value "LR", :left nil, :right nil}
 
 {:value "R",
  :left {:value "RL", :left nil, :right nil},
  :right {:value "RR", :left nil, :right nil}}
 
 {:value "RL", :left nil, :right nil}

 {:value "RR", :left nil, :right nil})

それぞれのトップ階層の :value を取り出せば、深さ優先で全ノードの :value を順回したことになります。

(map :value (tree-seq branch? children tree))
;;=> ("root" "L" "LL" "LR" "R" "RL" "RR")

このように通常は再帰を使う「木の深さ優先探索」を、シーケンス処理に置き換えてくれるのが tree-seq です。

◆階層化リストで木を作る

シンプルに次のように木構造を作ることもできます。

(def tree-list
     '("root"
       ("L"
        ("LL")
        ("LR"))
       ("R"
        ("RL")
        ("RR"))))

(def tree-list
     '("root"
       ("L" ("LL") ("LR"))
       ("R" ("RL") ("RR"))))

どちらも構造は同じです。
この tree-list も tree-seq で深さ優先探索することができます。

(tree-seq list? rest tree-list)

評価値は次のようなシーケンスです。一行が一要素になってます。

(("root" ("L" ("LL") ("LR")) ("R" ("RL") ("RR")))
 ("L" ("LL") ("LR"))
 ("LL")
 ("LR")
 ("R" ("RL") ("RR"))
 ("RL")
 ("RR"))

各要素の first をとれば tree-list の巡回になります。

(map first (tree-seq list? rest tree-list))
;;=> ("root" "L" "LL" "LR" "R" "RL" "RR")

tree-seq の引数は「list? が branch?」「rest が children」に対応しています。
tree-seq はこのように高階関数にすることで抽象化され、構造がどうであれ木であれば扱える関数になっています。

◆tree-seq と flatten

tree-list とその巡回結果を見比べると、list の平坦化が行われているともみなせます。ただし平坦化できるリストの構造には制限があります。限定的な flatten なので flatten- とします。

(defn flatten- [coll]
  (map first (tree-seq list? rest coll)))

(flatten- '("root" ("L" ("LL") ("LR")) ("R" ("RL") ("RR"))))
;;=> ("root" "L" "LL" "LR" "R" "RL" "RR")

(flatten- ((("LL") "L" ("LR")) "root" (("RL") ("RR") "R")))
;;=> java.lang.String cannot be cast to clojure.lang.IFn
;;   [Thrown class java.lang.ClassCastException]

後者の階層化リストも木の構造になっています。しかし、record で Node を作った場合とは異なり、このケースでは first や rest といったリストの要素の順序に依存した処理をしています。そのため後者のようなリストはうまく処理できずエラーになってしまいます。
ではどうすればいいのか？

◆clojure.core の flatten

考えるのサボって答えを見てしまいましょう。

(defn flatten
  "Takes any nested combination of sequential things (lists, vectors,
  etc.) and returns their contents as a single, flat sequence.
  (flatten nil) returns nil."
  {:added "1.2"}
  [x]
  (filter (complement sequential?)
          (rest (tree-seq sequential? seq x))))

今回のケースでは、sequential? は list? に置き換えても同じです。
まず tree-seq の部分からみてみると。

(def x '((("LL") "L" ("LR")) "root" (("RL") ("RR") "R")))
(tree-seq list? seq x)

結果

(
 ((("LL") "L" ("LR")) "root" (("RL") ("RR") "R"))
 (("LL") "L" ("LR"))
 ("LL")
 "LL"
 "L"
 ("LR")
 "LR"
 "root"
 (("RL") ("RR") "R")
 ("RL")
 "RL"
 ("RR")
 "RR"
 "R"
 )

children 関数を seq としているところがポイント。left,right の子ノードだけじゃなく、ノードが保持する値 (value) まで、子ノードとして扱っている。例えばルートノードの子ノードは便宜的に (左ノード "root" 右ノード）の３つとしています。
あとは tree-seq がリストを要素のレベルまで巡回してくれるので、得られたシーケンスのリストじゃない項を集めれば、平坦化されるというわけ。

(filter (complement list?) (tree-seq list? seq x))
;;=> ("LL" "L" "LR" "root" "RL" "RR" "R")

tree-seq の評価結果の最初の項は必ずリストになるので rest 除外しています。（どっちにしろ filter で除外されるから rest はなくてもいいと思うけど）。

最後に list? を sequential? に戻せば、clojure.core の flatten になります。

(defn flatten
  [x]
  (filter (complement sequential?)
          (rest (tree-seq sequential? seq x))))

sequential? はリストの他に、シーケンスやベクタに対しても true になります。
これで理解完了。

◆set や map は？

clojure の flatten では set や map は処理されませんん。

(flatten #{1 2 #{3} 4})     ;=> ()
(flatten {:a 1 :b {:c 10}}) ;=> ()
(flatten ['(1 2 (3)) #{4 5} 6 {:a 7}]) ;=> (1 2 3 #{4 5} 6 {:a 7})

でも、もうここまでくれば、set や map に対応した flatten も書けますね。

(defn flatten+
  [x]
  (filter (complement coll?)
          (rest (tree-seq coll? seq x))))

(flatten+ #{1 2 #{3} 4})     ;=> (1 2 3 4)

(flatten+ {:a 1 :b {:c 10}}) ;=> (:a 1 :b :c 10)

(flatten+ ['(1 2 (3)) #{4 5} 6 {:a 7}]) ;=> (1 2 3 4 5 6 :a 7)

(flatten+ {:a 1 :b #{10 20 #{30} #{:x :y} :c {:d 2 :e 3}}})
;;=> (:a 1 :b :d 2 :e 3 10 :c :y :x 20 30)

set や map は一般に順序が保存されないので、flatten+ の結果は順序が保証されません。
コレクションと sequential? coll? の関係は以前調べてまとめたものがあります。参考までに。
=> seq? sequential? coll?

以上