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:

-
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)))))))

-
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)

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:

https://github.com/clojure/clojure/blob/2224dbad5534ff23d3653b07d9dc0a60ba076dd7/src/clj/clojure/walk.clj#L35

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).

StackExchange Code Review Q#113703, answer score: 2