Skip to content

Instantly share code, notes, and snippets.

@niwinz
Forked from stathissideris/tree-seq-extra.clj
Created June 11, 2020 06:26
Show Gist options
  • Save niwinz/129899889ea1c5a241c5a4a84dc54230 to your computer and use it in GitHub Desktop.
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)
(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