Postorder tree traversal with clojure.zip to edit nodes

Question

I have a tree represented as a nested vector. I want to have a generalization of indexed for trees, showing the index of each node like this,

(visit 42); => [0 42]
(visit [6 7]); => [0
              ;     [[0 6] 
              ;      [1 7]]]

The naive implementation would use clojure.zip directly (as already asked here)

(defn visit [tree]
  (loop [loc (vector-zip tree)]
    (if (end? loc)
      (root loc)
      (recur 
        (next (edit loc #(conj 
                           [(count (lefts loc))] 
                           %)))))))

But recurring with clojure.zip/next performs a preorder traversal, resulting in an infinite loop in this case (unvisited nodes get conjed into a [:found] vector infinitely). Another approach would be using clojure.walk/postwalk, but it doesn't provide structural information, such as index.

How would you implement this? Is there a postorder-next for zip that would solve it right away?

Answer 1

I'm not sure if I understand what you're trying to do, but the examples suggest to me that the indices assigned to the nodes correspond to the number of their left siblings (since in the second example both the root node and the 6 child are labelled with 0). Update: Apparently I simply misread the visit example the first time round -- it makes the intention clear enough... Fortunately now that I read it properly it seems to me that the answer below is correct.

If that is correct, this is a clojure.walk/postwalk-based solution:

(defn index-vec-tree [vt]
  [0 (walk/postwalk
       (fn [node]
         (if-not (vector? node)
           node
           (vec (map-indexed vector node))))
       vt)])

With the given examples:

user> (index-vec-tree [6 7])
[0 [[0 6] [1 7]]]
user> (index-vec-tree 42)
[0 42]

Update: A simple solution using map-indexed (available in 1.2; use map + clojure.contrib.seq-utils/indexed in 1.1):

(defn index-vec-tree [vt]
  (letfn [(iter [i vt] [i (if (vector? vt)
                                (vec (map-indexed iter vt))
                                vt)])]
    (iter 0 vt)))