Squaring a tree in Clojure

I am working on Problem 2.30 from Structure and Interpretation of Computer Programs. I book is in scheme, but I am doing the exercises in Clojure.

The problem is to write code that takes a tree of numbers and return a new tree with the numbers squared. I have done it in two ways:

1. Without higher-order functions:

(defn square-tree1 [tree]
(lazy-seq
(when-let [s (seq tree)]
(if (coll? (first s))
(cons (square-tree1 (first s))
(square-tree1 (rest s)))
(cons (square (first s))
(square-tree1 (rest s)))))))

2. Using map

(defn square-tree2 [tree]
(if (coll? tree)
(map square-tree2 tree)
(square tree)))


Where:

(defn square [x]
(* x x))


My tests for this are:

(deftest e2.30a
(testing "Ex 2.30a: square-tree"
(is (=
(square-tree1
(list 1
(list 2 (list 3 4) 5)
(list 6 7)))
'(1 (4 (9 16) 25) (36 49))))
(is (= (square-tree1 '()) '()))
(is (= (square-tree1 [1 2 3 4 5]) [1 4 9 16 25]))
(is (= (square-tree1 [1 [2 [3]]]) [1 [4 [9]]]))))

• I want the functions to work with lists and vector inputs.
• I am unsure if 1 is the proper usage of lazy-seq and whether there is a neater way to express this.
• Is (col? (first s)) the correct predicate here? I find it tricky in Clojure to test if the node in the tree is another sub-tree or if it is number. In Scheme it is easier because you have the pair? predicate. I would like to see if there is a better way.

Edit -- After a useful discussion with Timothy Pratley about clojure.walk

After looking at the source code of clojure.walk I have come up with a third square-tree function that extends the square-tree2, but now preserves the input collection types of the input tree like clojure.walk does:

    (defn square-tree3
[tree]
(cond
(list? tree) (apply list (map square-tree3 tree))
(seq? tree) (doall (map square-tree3 tree))
(coll? tree) (into (empty tree) (map square-tree3 tree))
:else (square tree)))


So it extends the answer to the coll? predicate with a (into (empty tree)...) which puts the result of map into whatever type of collection tree is. It also adds cases for list and seq types, handling each of them in their own special way so that they replicate into the same collection type.

You can use postwalk:

(defn maybe-square [x]
(if (number? x)
(* x x)
x))

(clojure.walk/postwalk maybe-square [1 [2 [3]]])
=> [1 [4 [9]]]


Which is perhaps against the spirit of learning how to do it, so it is worth looking at the implementation code:

Which shows the correct way to test for various things you may expect in a tree (There are quite a few different tests you can take advantage of).

lazy-seq and map are cool, but it is probably more common to think of trees as data structures rather than lazy sequences; you will note that postwalk returns datastructures, not lazy sequences. So I recommend you also return datastructures (in your case vectors and lists, rather than a sequence).

• Thank you. In the wild I'd use zippers/tree-seq/walk for handling trees. Here I'm trying to do things explicitly for my own education. However the datastructure vs seq thing has confused me about Clojure. Are there some best practices for when a function should return a vector and when to return a seq? It feels with Clojure like everything starts with a vector, then I do a map or a filter and then everything becomes a seq and I've lost my vector. There is no map that retains data type? – jimypbr Dec 13 '15 at 18:16
• Seqs are used extensively in Clojure programs to great effect, but clearly there are situations where a vector is better than a seq, because of constant time lookup (calling nth on a seq is O(n), calling nth on a vector is O(1)). So the answer to your question "when should a function return a vector instead of a seq?" is when it is convenient to use that result in a way where O(1) access is important. As for "map that retains data type" Take a look at vec, mapv, into and empty. github.com/nathanmarz/specter is a library that has an explicit goal of preserving the input type. – Timothy Pratley Dec 13 '15 at 18:53
• I recommend reading the source code of clojure.walk because it addresses these concerns that you raise about 'preserving types' in an idiomatic fashion. The code presents several design decisions relevant to the problem you are studying. So there is much to learn from it. Seeing you already wrote a solution on your own, you will be able to appreciate the tradeoffs, and identify why they might be good or bad. – Timothy Pratley Dec 13 '15 at 19:16
• That was a useful exercise, thanks. I've added a new version of square-tree inspired by clojure.walk. – jimypbr Dec 15 '15 at 3:06

Starting with your square-tree2, we can distinguish vectors from other sequences as follows:

(defn square-tree2 [tree]
(if (coll? tree)
((if (vector? tree) vec identity) (map square-tree2 tree))
(square tree)))

• The map makes the horizontal explorations lazy.
• An enclosing lazy-seq would not make vertical exploration lazy - each level still involves a recursive call.

If you want to roll your own map, as square-tree1 does, you can simplify things by pulling it out into a distinct function:

(defn map-square [s]
(lazy-seq
(when (seq s)
(cons (square-tree (first s)) (map-square (rest s))))))


then

(defn square-tree [tree]
(if (coll? tree)
(map-square tree)
(square tree)))


These are mutually recursive, so we need to (declare square-tree) up front.

Edited to correct error in map-square.