Squaring a tree in Clojure

Question

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.

Answer 1

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

Answer 2

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.