-
-
Save niwinz/129899889ea1c5a241c5a4a84dc54230 to your computer and use it in GitHub Desktop.
Like Clojure's tree-seq, but with depth info for each node or the full path (recursive - blow up the stack for deep trees)
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
(defn tree-seq-depth | |
"Returns a lazy sequence of vectors of the nodes in a tree and their | |
depth as [node depth], 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." | |
[branch? children root] | |
(let [walk (fn walk [depth node] | |
(lazy-seq | |
(cons [node depth] | |
(when (branch? node) | |
(mapcat (partial walk (inc depth)) (children node))))))] | |
(walk 0 root))) | |
(defn tree-seq-path | |
"Like core's tree-seq but returns a lazy sequence of vectors of the | |
paths of the nodes in a tree, via a depth-first walk. It optionally | |
applies node-fn to each node before adding it to the path. 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." | |
[branch? children root & [node-fn]] | |
(let [node-fn (or node-fn identity) | |
walk (fn walk [path node] | |
(let [new-path (conj path (node-fn node))] | |
(lazy-seq | |
(cons new-path | |
(when (branch? node) | |
(mapcat (partial walk new-path) (children node)))))))] | |
(walk [] root))) | |
;;try with: | |
(def tree | |
{:name "foo" | |
:children | |
[{:name "bar1"} | |
{:name "bar2" | |
:children | |
[{:name "baz1"} | |
{:name "baz2"} | |
{:name "baz3"} | |
{:name "baz4"} | |
{:name "baz5"}]} | |
{:name "bar3"} | |
{:name "bar4" | |
:children | |
[{:name "biz1"} | |
{:name "biz2"} | |
{:name "biz3"}]}]}) | |
(deftest test-tree-seq-path | |
(= | |
[["foo"] | |
["foo" "bar1"] | |
["foo" "bar2"] | |
["foo" "bar2" "baz1"] | |
["foo" "bar2" "baz2"] | |
["foo" "bar2" "baz3"] | |
["foo" "bar2" "baz4"] | |
["foo" "bar2" "baz5"] | |
["foo" "bar3"] | |
["foo" "bar4"] | |
["foo" "bar4" "biz1"] | |
["foo" "bar4" "biz2"] | |
["foo" "bar4" "biz3"]] | |
(tree-seq-path :children :children tree :name))) | |
(deftest test-tree-seq-depth | |
(is (= | |
[["foo" 0] | |
["bar1" 1] | |
["bar2" 1] | |
["baz1" 2] | |
["baz2" 2] | |
["baz3" 2] | |
["baz4" 2] | |
["baz5" 2] | |
["bar3" 1] | |
["bar4" 1] | |
["biz1" 2] | |
["biz2" 2] | |
["biz3" 2]] | |
(map #(update-in % [0] :name) | |
(tree-seq-depth :children :children tree))))) |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment