(defn tree-map [mapper branch? children root]
  (let [r (lazy-seq (map (fn zmap-mapper [e] (if (branch? e) (tree-map mapper branch? children (children e)) (mapper e))) root))]
    (cond 
     (vector? root)
     (vec r)
     (associative? root)
     (reduce (fn [m [k v]] (assoc m k v)) {} r)
     (set? root)
     (set r)
     :else 
     r)))